@@ -123,10 +123,7 @@ class Granite(SimpleService): **input_message.info, } ) - try: - result = ''.join(self._call_api(payload=callback_data)['output']) - except Exception: - raise ExternalAPIException + result = ''.join(self._call_api(payload=callback_data)['output']) process_time = timedelta(seconds=(time.time() - start_time)) self.handle_invoice(input_message.content_object.model, result, input_message) msgs = self.save_results(result, process_time, save) @@ -140,10 +140,7 @@ class Iconic(SimpleService): if input_message.file: callback_data.update({'image': BytesIO(input_message.file.read())}) input_message.file.close() - try: - images = replicate_run(self._CALLBACK, callback_data) - except Exception: - raise ExternalAPIException() + images = replicate_run(self._CALLBACK, callback_data) process_time = timedelta(seconds=(time.time() - start_time)) self.handle_invoice(input_message.content_object.model, process_time=process_time) msgs = self.save_results(input_message.content, images, process_time, save) @@ -119,10 +119,7 @@ class Lightning(SimpleService): **input_message.info, } ) - try: - images = replicate_run(self._CALLBACK, callback_data) - except Exception: - raise ExternalAPIException() + images = replicate_run(self._CALLBACK, callback_data) process_time = timedelta(seconds=(time.time() - start_time)) self.handle_invoice(input_message.content_object.model, process_time=process_time) msgs = self.save_results(input_message.content, images, process_time, save) @@ -123,12 +123,7 @@ class Logoai(SimpleService): if input_message.file: callback_data.update({'image': BytesIO(input_message.file.read())}) input_message.file.close() - try: - images = replicate_run(self._CALLBACK, callback_data) - except ReplicateError: - raise ExternalAPIException() - except ModelError as e: - raise InvalidDataException() + images = replicate_run(self._CALLBACK, callback_data) process_time = timedelta(seconds=(time.time() - start_time)) self.handle_invoice(input_message.content_object.model, process_time=process_time) msgs = self.save_results(input_message.content, images, process_time, save) @@ -141,10 +141,7 @@ class Pulid(SimpleService): if input_message.file: callback_data.update({'main_face_image': BytesIO(input_message.file.read())}) input_message.file.close() - try: - images = replicate_run(self._CALLBACK, callback_data) - except Exception: - raise ExternalAPIException() + images = replicate_run(self._CALLBACK, callback_data) process_time = timedelta(seconds=(time.time() - start_time)) self.handle_invoice(input_message.content_object.model, process_time=process_time) msgs = self.save_results(input_message.content, images, process_time, save) @@ -116,27 +116,24 @@ class Recraft(SimpleService): return f'{size[0]}x{size[1]}' def _call_api(self, payload: dict) -> list: - try: - headers = { - 'Authorization': f'Bearer {settings.REPLICATE_API_KEY}', - 'Content-Type': 'application/json', - 'Prefer': 'wait', - } - data = {'input': payload} - response = requests.post( - url=f'{self.generate_url}{payload['version']}/predictions', - headers=headers, - json=data, - ) - if response.status_code != 201: - raise Exception(response.json()) + headers = { + 'Authorization': f'Bearer {settings.REPLICATE_API_KEY}', + 'Content-Type': 'application/json', + 'Prefer': 'wait', + } + data = {'input': payload} + response = requests.post( + url=f'{self.generate_url}{payload['version']}/predictions', + headers=headers, + json=data, + ) + if response.status_code != 201: + raise Exception(response.json()) + result = requests.get(url=f'{self.get_url}{response.json().get('id')}', headers=headers) + while result.json()['status'] not in ('succeeded', 'failed', 'canceled'): result = requests.get(url=f'{self.get_url}{response.json().get('id')}', headers=headers) - while result.json()['status'] not in ('succeeded', 'failed', 'canceled'): - result = requests.get(url=f'{self.get_url}{response.json().get('id')}', headers=headers) - return result.json()['output'] - except Exception: - raise ExternalAPIException + return result.json()['output'] def calculate_price(self, input_message: Message) -> Decimal: return self.payment_rules[input_message.info.get('version', 'recraft-v3')] @@ -121,10 +121,7 @@ class Sdxlemoji(SimpleService): if input_message.file: callback_data.update({'image': BytesIO(input_message.file.read())}) input_message.file.close() - try: - images = replicate_run(self._CALLBACK, callback_data) - except Exception: - raise ExternalAPIException() + images = replicate_run(self._CALLBACK, callback_data) process_time = timedelta(seconds=(time.time() - start_time)) self.handle_invoice(input_message.content_object.model, process_time=process_time) msgs = self.save_results(input_message.content, images, process_time, save)