@@ -6,7 +6,6 @@ from typing import Any import requests from django.core.files import File -from replicate.exceptions import ModelError from messages.models import Message from ml_model.exceptions import RequestBlocked, GenerationException @@ -58,12 +57,7 @@ class Dalle(SimpleService): **input_message.info, } ) - try: - images = replicate_run(self._CALLBACK, callback_data) - except ModelError as exc: - if any(error in str(exc) for error in ('E005', 'E006', 'sexual', 'NSFW')): - raise RequestBlocked - raise GenerationException from exc + images = replicate_run(self._CALLBACK, callback_data) process_time = timedelta(seconds=(time.time() - start_time)) self.handle_invoice(input_message.content_object.model, input_message=input_message) msgs = self.save_results(input_message.content, images, process_time, save) @@ -7,7 +7,6 @@ from typing import Any import requests from django.core.files import File from django.utils.translation import gettext as _ -from replicate.exceptions import ModelError from messages.models import Message from ml_model.exceptions import RequestBlocked, GenerationException, InvalidParameterError @@ -54,12 +53,7 @@ class Elevenlabs_Music(SimpleService): **input_message.info, } start_time = time.time() - try: - audio = replicate_run('elevenlabs/music', callback_data) - except ModelError as exc: - if any(error in str(exc) for error in ('E005', 'E006', 'sexual')): - raise RequestBlocked - raise GenerationException from exc + audio = replicate_run('elevenlabs/music', callback_data) process_time = timedelta(seconds=(time.time() - start_time)) self.handle_invoice(input_message.content_object.model, duration=duration) msgs = self.save_results(input_message.content, process_time, audio, save) @@ -10,7 +10,6 @@ import requests from PIL import Image from django.core.files import File from django.core.files.images import get_image_dimensions -from replicate.exceptions import ModelError from messages.models import Message from ml_model.exceptions import PredictionInterruptedError, RequestBlocked, GenerationException, \ @@ -96,14 +95,7 @@ class Flux_2(SimpleService): image = f'data:image/{format};base64,{base64.b64encode(buf.getvalue()).decode("utf-8")}' buf.close() callback_data.update({'input_images': [image]}) - try: - images = [replicate_run(f'black-forest-labs/{version}', callback_data)] - except ModelError as exc: - if any(error in str(exc) for error in ('E005', 'E006', 'sexual')): - raise RequestBlocked - elif 'PA' in str(exc): - raise PredictionInterruptedError - raise GenerationException from exc + images = [replicate_run(f'black-forest-labs/{version}', callback_data)] process_time = timedelta(seconds=(time.time() - start_time)) self.handle_invoice(input_message.content_object.model, version=version, input_mp=input_mp, output_mp=output_mp) msgs = self.save_results(input_message.content, images, process_time, save) @@ -5,7 +5,6 @@ import time from typing import Any from django.core.files import File -from replicate.exceptions import ModelError import requests from messages.models.message import Message @@ -73,15 +72,7 @@ class Fluxpulid(SimpleService): } ) start_time = time.time() - try: - images = replicate_run(self.ENDPOINT, callback_data) - except ModelError as exc: - if any(error in str(exc) for error in ('E005', 'E006', 'sexual')): - raise RequestBlocked - elif exc.prediction.error == 'facexlib align face fail': - raise FaceNotFoundError - raise GenerationException from exc - + images = replicate_run(self.ENDPOINT, callback_data) process_time = timedelta(seconds=(time.time() - start_time)) self.handle_invoice(input_message.content_object.model, len(images)) msgs = self.save_results(input_message.content, images, process_time, save) @@ -6,7 +6,6 @@ from typing import Any import requests from django.core.files import File -from replicate.exceptions import ModelError from messages.models import Message from ml_model.exceptions import ModelTimeoutError, ImageContentNotFound, RequestBlocked, GenerationException @@ -44,25 +43,15 @@ class Geminiimage(SimpleService): def make(self, input_message: Message, save: bool = True) -> list[Message]: if input_message.content: - try: - callback_data = dict( - {'prompt': self.translate_prompt(input_message.content), **input_message.info} - ) - start_time = time.time() - images = replicate_run('google/gemini-2.5-flash-image', callback_data) - process_time = timedelta(seconds=(time.time() - start_time)) - self.handle_invoice( - input_message.content_object.model, num_images=input_message.info.get('num_images', 1) - ) - msgs = self.save_results(input_message.content, process_time, images, save) - return msgs - except ModelError as exc: - if exc.prediction.error in ( - 'No image content found in response', - 'Failed to generate image.', - ): - raise ImageContentNotFound - elif any(error in str(exc) for error in ('E005', 'E006', 'sexual')): - raise RequestBlocked - raise GenerationException from exc + callback_data = dict( + {'prompt': self.translate_prompt(input_message.content), **input_message.info} + ) + start_time = time.time() + images = replicate_run('google/gemini-2.5-flash-image', callback_data) + process_time = timedelta(seconds=(time.time() - start_time)) + self.handle_invoice( + input_message.content_object.model, num_images=input_message.info.get('num_images', 1) + ) + msgs = self.save_results(input_message.content, process_time, images, save) + return msgs raise ModelTimeoutError @@ -7,7 +7,6 @@ from io import BytesIO from typing import Any from django.core.files import File -from replicate.exceptions import ModelError from messages.models import Message from ml_model.exceptions import ( @@ -54,14 +53,7 @@ class Grok_Image(SimpleService): if image := input_message.file: callback_data.update({'image': image.url}) start_time = time.time() - try: - images = replicate_run('xai/grok-imagine-image', callback_data) - except ModelError as exc: - if any(error in str(exc) for error in ('E005', 'E006', 'sexual', 'NSFW')): - raise RequestBlocked - elif exc.prediction.error == 'No image content found in response': - raise ImageContentNotFound - raise GenerationException from exc + images = replicate_run('xai/grok-imagine-image', callback_data) process_time = timedelta(seconds=(time.time() - start_time)) self.handle_invoice(input_message.content_object.model) msgs = self.save_results(input_message.content, process_time, images, save) @@ -8,7 +8,6 @@ from typing import Any import filetype import requests from django.core.files import File -from replicate.exceptions import ModelError from messages.models import Message from ml_model.exceptions import RequestBlocked, GenerationException @@ -52,12 +51,7 @@ class Grok_Imagine_Video(SimpleService): if image := input_message.file: callback_data.update({'image': image.url}) start_time = time.time() - try: - video = replicate_run('xai/grok-imagine-video', callback_data) - except ModelError as exc: - if any(error in str(exc) for error in ('E005', 'E006', 'sexual')): - raise RequestBlocked - raise GenerationException from exc + video = replicate_run('xai/grok-imagine-video', callback_data) process_time = timedelta(seconds=(time.time() - start_time)) self.handle_invoice(input_message.content_object.model, duration=duration) msgs = self.save_results(input_message.content, process_time, video, save) @@ -8,7 +8,6 @@ from typing import Any import filetype import requests from django.core.files import File -from replicate.exceptions import ModelError from messages.models import Message from ml_model.exceptions import RequestBlocked, GenerationException @@ -70,15 +69,10 @@ class Hailuo(SimpleService): input_message.file.close() callback_data.update({'first_frame_image': image}) start_time = time.time() - try: - video = replicate_run( - f'minimax/{version}', - callback_data - ) - except ModelError as exc: - if any(error in str(exc) for error in ('E005', 'E006', 'sexual')): - raise RequestBlocked - raise GenerationException from exc + video = replicate_run( + f'minimax/{version}', + callback_data + ) process_time = timedelta(seconds=(time.time() - start_time)) self.handle_invoice(input_message.content_object.model, version=version, resolution=resolution) msgs = self.save_results(input_message.content, process_time, video, save) @@ -6,7 +6,6 @@ from typing import Any import requests from django.core.files import File -from replicate.exceptions import ModelError from messages.models import Message from ml_model.exceptions import InvalidStyleCombinationError, RequestBlocked, GenerationException @@ -80,15 +79,10 @@ class Ideogram(SimpleService): **input_message.info, } ) - try: - runner = replicate_run( - f'{self._CALLBACK_BASE}{callback_data.get("version", "ideogram-v3-turbo")}', - callback_data, - ) - except ModelError as exc: - if any(error in str(exc) for error in ('E005', 'E006', 'sexual')): - raise RequestBlocked - raise GenerationException from exc + runner = replicate_run( + f'{self._CALLBACK_BASE}{callback_data.get("version", "ideogram-v3-turbo")}', + callback_data, + ) images = runner if isinstance(runner, list) else [runner] process_time = timedelta(seconds=(time.time() - start_time)) self.handle_invoice(input_message.content_object.model, input_message=input_message, version=version) @@ -6,7 +6,6 @@ from typing import Any import requests from django.core.files import File -from replicate.exceptions import ModelError from messages.models import Message from ml_model.exceptions import ImageContentNotFound, GenerationException, RequestBlocked @@ -38,25 +37,18 @@ class Imagen(SimpleService): return [msg] def make(self, input_message: Message, save: bool = True) -> list[Message]: - try: - if (balance := PaymentPlanSelector(self.store.user).get_current_balance()) < self.TOKENS_COST: - raise InsufficientBalance(balance, self.TOKENS_COST) - callback_data = dict( - { - 'prompt': self.translate_prompt(input_message.content), - 'safety_filter_level': 'block_medium_and_above', - **input_message.info, - } - ) - start_time = time.time() - images = replicate_run('google/imagen-3-fast', callback_data) - process_time = timedelta(seconds=(time.time() - start_time)) - self.handle_invoice(input_message.content_object.model) - msgs = self.save_results(input_message.content, process_time, images, save) - return msgs - except ModelError as exc: - if any(error in str(exc) for error in ('E005', 'E006', 'sexual', 'NSFW')): - raise RequestBlocked - elif exc.prediction.error == 'No image content found in response': - raise ImageContentNotFound - raise GenerationException from exc + if (balance := PaymentPlanSelector(self.store.user).get_current_balance()) < self.TOKENS_COST: + raise InsufficientBalance(balance, self.TOKENS_COST) + callback_data = dict( + { + 'prompt': self.translate_prompt(input_message.content), + 'safety_filter_level': 'block_medium_and_above', + **input_message.info, + } + ) + start_time = time.time() + images = replicate_run('google/imagen-3-fast', callback_data) + process_time = timedelta(seconds=(time.time() - start_time)) + self.handle_invoice(input_message.content_object.model) + msgs = self.save_results(input_message.content, process_time, images, save) + return msgs @@ -9,7 +9,6 @@ import filetype import requests from django.core.files import File -from replicate.exceptions import ModelError from messages.models import Message from ml_model.exceptions import RequestBlocked, GenerationException, FileNotProvided @@ -60,12 +59,7 @@ class Kling(SimpleService): input_message.file.close() callback_data.update({'start_image': image}) start_time = time.time() - try: - video = replicate_run('kwaivgi/kling-v2.1', callback_data) - except ModelError as exc: - if any(error in str(exc) for error in ('E005', 'E006', 'sexual')): - raise RequestBlocked - raise GenerationException from exc + video = replicate_run('kwaivgi/kling-v2.1', callback_data) process_time = timedelta(seconds=(time.time() - start_time)) self.handle_invoice(input_message.content_object.model, mode=mode, duration=duration) msgs = self.save_results(input_message.content, process_time, video, save) @@ -6,7 +6,6 @@ from typing import Any import requests from django.core.files import File -from replicate.exceptions import ModelError from messages.models import Message from ml_model.exceptions import RequestBlocked, GenerationException @@ -73,15 +72,10 @@ class Leonardo(SimpleService): **input_message.info, } ) - try: - runner = replicate_run( - f'{self._CALLBACK_BASE}{callback_data.get("version", "lucid-origin")}', - callback_data, - ) - except ModelError as exc: - if any(error in str(exc) for error in ('E005', 'E006', 'sexual')): - raise RequestBlocked - raise GenerationException from exc + runner = replicate_run( + f'{self._CALLBACK_BASE}{callback_data.get("version", "lucid-origin")}', + callback_data, + ) images = runner if isinstance(runner, list) else [runner] process_time = timedelta(seconds=(time.time() - start_time)) self.handle_invoice( @@ -6,7 +6,6 @@ from typing import Any import requests from django.core.files import File -from replicate.exceptions import ModelError from messages.models import Message from ml_model.exceptions import RequestBlocked, GenerationException @@ -47,12 +46,7 @@ class Lyria(SimpleService): if file := input_message.file: callback_data.update({'images': [file.url]}) start_time = time.time() - try: - video = replicate_run(f'google/{version}', callback_data) - except ModelError as exc: - if any(error in str(exc) for error in ('E005', 'E006', 'sexual')): - raise RequestBlocked - raise GenerationException from exc + video = replicate_run(f'google/{version}', callback_data) process_time = timedelta(seconds=(time.time() - start_time)) self.handle_invoice(input_message.content_object.model, version) msgs = self.save_results(input_message.content, process_time, video, save) @@ -7,7 +7,6 @@ from typing import Any import requests from django.core.files import File from django.utils.translation import gettext as _ -from replicate.exceptions import ModelError from messages.models import Message from ml_model.exceptions import RequestBlocked, InvalidParameterError, GenerationException @@ -51,14 +50,7 @@ class Minimaxmusic(SimpleService): **input_message.info, } start_time = time.time() - try: - audio = replicate_run('minimax/music-1.5', callback_data) - except ModelError as exc: - if 'lyrics is too long' in str(exc): - raise InvalidParameterError(_('Lyrics is too long')) - if any(error in str(exc) for error in ('E005', 'E006', 'sexual')): - raise RequestBlocked - raise GenerationException from exc + audio = replicate_run('minimax/music-1.5', callback_data) process_time = timedelta(seconds=(time.time() - start_time)) self.handle_invoice(input_message.content_object.model) msgs = self.save_results(input_message.content, process_time, audio, save) @@ -7,7 +7,6 @@ from typing import Any import requests from django.core.files import File from django.utils.translation import gettext as _ -from replicate.exceptions import ModelError from messages.models import Message from ml_model.exceptions import RequestBlocked, InvalidParameterError, GenerationException @@ -55,14 +54,7 @@ class Minimaxmusic_Lite(SimpleService): file_url = Preset.objects.get(slug=instrumental).file.url callback_data.update({'instrumental_file': file_url}) start_time = time.time() - try: - audio = replicate_run('minimax/music-01', callback_data) - except ModelError as exc: - if 'lyrics is too long' in str(exc): - raise InvalidParameterError(_('Lyrics is too long')) - if any(error in str(exc) for error in ('E005', 'E006', 'sexual')): - raise RequestBlocked - raise GenerationException from exc + audio = replicate_run('minimax/music-01', callback_data) process_time = timedelta(seconds=(time.time() - start_time)) self.handle_invoice(input_message.content_object.model) msgs = self.save_results(input_message.content, process_time, audio, save) @@ -8,7 +8,6 @@ from typing import Any import filetype import requests from django.core.files import File -from replicate.exceptions import ModelError from messages.models import Message from ml_model.exceptions import RequestBlocked, GenerationException @@ -55,12 +54,7 @@ class Minimaxvideo(SimpleService): input_message.file.close() callback_data.update({'subject_reference': image}) start_time = time.time() - try: - video = replicate_run(f'minimax/{version}', callback_data) - except ModelError as exc: - if any(error in str(exc) for error in ('E005', 'E006', 'sexual')): - raise RequestBlocked - raise GenerationException from exc + video = replicate_run(f'minimax/{version}', callback_data) process_time = timedelta(seconds=(time.time() - start_time)) self.handle_invoice(input_message.content_object.model, version) msgs = self.save_results(input_message.content, process_time, video, save) @@ -8,7 +8,6 @@ from typing import Optional, Any import filetype import requests from django.core.files import File -from replicate.exceptions import ModelError from messages.models import Message from ml_model.exceptions import ( @@ -78,16 +77,7 @@ class Nanobanana(SimpleService): image = f'data:{mime};base64,{base64.b64encode(input_message.file.read()).decode("utf-8")}' input_message.file.close() callback_data.update({'image_input': [image]}) - try: - image = replicate_run(f'google/{version}', callback_data) - except ModelError as exc: - if exc.prediction.error == 'No image content found in response': - raise ImageContentNotFound - elif exc.prediction.error in ('400', 'Failed to generate image.'): - raise ModelCouldNotInterpretPrompt from exc - elif any(error in str(exc) for error in ('E005', 'E006', 'sexual')): - raise RequestBlocked - raise GenerationException from exc + image = replicate_run(f'google/{version}', callback_data) process_time = timedelta(seconds=(time.time() - start_time)) self.handle_invoice(input_message.content_object.model, version=version, resolution=resolution) msgs = self.save_results(input_message.content, image, process_time, save) @@ -8,7 +8,6 @@ from typing import Any, Optional import filetype import requests from django.core.files import File -from replicate.exceptions import ModelError from messages.models import Message from ml_model.exceptions import ( @@ -77,16 +76,7 @@ class Nanobanana_2(SimpleService): image = f'data:{mime};base64,{base64.b64encode(input_message.file.read()).decode("utf-8")}' input_message.file.close() callback_data.update({'image_input': [image]}) - try: - image = replicate_run(f'google/{version}', callback_data) - except ModelError as exc: - if exc.prediction.error == 'No image content found in response': - raise ImageContentNotFound - elif exc.prediction.error in ('400', 'Failed to generate image.'): - raise ModelCouldNotInterpretPrompt from exc - elif any(error in str(exc) for error in ('E005', 'E006', 'sexual')): - raise RequestBlocked - raise GenerationException from exc + image = replicate_run(f'google/{version}', callback_data) process_time = timedelta(seconds=(time.time() - start_time)) self.handle_invoice(input_message.content_object.model, version=version, resolution=resolution) msgs = self.save_results(input_message.content, image, process_time, save) @@ -10,7 +10,6 @@ import filetype import requests from django.core.files import File -from replicate.exceptions import ModelError from messages.models import Message from ml_model.exceptions import ( @@ -65,14 +64,7 @@ class Photon(SimpleService): input_message.file.close() callback_data.update({'image_reference': image}) start_time = time.time() - try: - images = replicate_run('luma/photon-flash', callback_data) - except ModelError as exc: - if any(error in str(exc) for error in ('E005', 'E006', 'sexual', 'NSFW')): - raise RequestBlocked - elif exc.prediction.error == 'No image content found in response': - raise ImageContentNotFound - raise GenerationException from exc + images = replicate_run('luma/photon-flash', callback_data) process_time = timedelta(seconds=(time.time() - start_time)) self.handle_invoice(input_message.content_object.model) msgs = self.save_results(input_message.content, process_time, images, save) @@ -7,7 +7,6 @@ from io import BytesIO from typing import Any from django.core.files import File -from replicate.exceptions import ModelError from messages.models import Message from ml_model.exceptions import RequestBlocked, GenerationException @@ -63,12 +62,7 @@ class Pixverse(SimpleService): if image := input_message.file: callback_data.update({'image': image.url}) start_time = time.time() - try: - video = replicate_run('pixverse/pixverse-v5.6', callback_data) - except ModelError as exc: - if any(error in str(exc) for error in ('E005', 'E006', 'sexual')): - raise RequestBlocked - raise GenerationException from exc + video = replicate_run('pixverse/pixverse-v5.6', callback_data) process_time = timedelta(seconds=(time.time() - start_time)) self.handle_invoice(input_message.content_object.model, quality=quality, duration=duration) msgs = self.save_results(input_message.content, process_time, video, save) @@ -9,7 +9,6 @@ import filetype import requests from django.core.files import File -from replicate.exceptions import ModelError from messages.models import Message from ml_model.exceptions import GenerationException, RequestBlocked @@ -74,14 +73,9 @@ class Pruna_V(SimpleService): else: callback_data.update({'image': input_message.file.url}) start_time = time.time() - try: - images = replicate_run('prunaai/p-video', callback_data) - if 'nsfw.jpeg' == str(images).split('/')[-1]: - raise RequestBlocked - except ModelError as exc: - if any(error in str(exc) for error in ('E005', 'E006', 'sexual', 'NSFW')): - raise RequestBlocked - raise GenerationException from exc + images = replicate_run('prunaai/p-video', callback_data) + if 'nsfw.jpeg' == str(images).split('/')[-1]: + raise RequestBlocked process_time = timedelta(seconds=(time.time() - start_time)) self.handle_invoice( input_message.content_object.model, @@ -8,7 +8,6 @@ from typing import Any import requests from django.core.files import File -from replicate.exceptions import ModelError from messages.models import Message from ml_model.exceptions import GenerationException, RequestBlocked @@ -70,12 +69,7 @@ class Prunaai(SimpleService): } callback_data.update({'speed_mode': speed_mode[s_m]}) start_time = time.time() - try: - images = replicate_run(f'prunaai/{version}', callback_data) - except ModelError as exc: - if any(error in str(exc) for error in ('E005', 'E006', 'sexual', 'NSFW')): - raise RequestBlocked - raise GenerationException from exc + images = replicate_run(f'prunaai/{version}', callback_data) process_time = timedelta(seconds=(time.time() - start_time)) self.handle_invoice(input_message.content_object.model, version=version) msgs = self.save_results(input_message.content, process_time, images, save) @@ -8,7 +8,6 @@ from typing import Any import filetype import requests from django.core.files import File -from replicate.exceptions import ModelError from messages.models import Message from ml_model.exceptions import RequestBlocked, GenerationException @@ -59,12 +58,7 @@ class Ray(SimpleService): input_message.file.close() callback_data.update({'start_image': image}) start_time = time.time() - try: - video = replicate_run(f'luma/{version}', callback_data) - except ModelError as exc: - if any(error in str(exc) for error in ('E005', 'E006', 'sexual')): - raise RequestBlocked - raise GenerationException from exc + video = replicate_run(f'luma/{version}', callback_data) process_time = timedelta(seconds=(time.time() - start_time)) self.handle_invoice(input_message.content_object.model, version=version, duration=duration) msgs = self.save_results(input_message.content, process_time, video, save) @@ -5,7 +5,6 @@ from decimal import Decimal from io import BytesIO from typing import Any -from replicate.exceptions import ModelError from ml_model.exceptions import ( RequestBlocked, GenerationException, @@ -71,16 +70,7 @@ class Reve(SimpleService): input_message.file.close() callback_data.update({'image': image}) type = 'edit-fast' - try: - image = replicate_run(f'reve/{type}', callback_data) - except ModelError as exc: - if any(error in str(exc) for error in ('E005', 'E006', 'sexual')): - raise RequestBlocked - if 'INPUT_ANALYSIS_FAILURE' in str(exc): - raise ImageAnalysisError - if 'PROMPT_TOO_LONG' in str(exc): - raise ExceededContextLengthError - raise GenerationException from exc + image = replicate_run(f'reve/{type}', callback_data) process_time = timedelta(seconds=(time.time() - start_time)) self.handle_invoice(input_message.content_object.model, type=type) msgs = self.save_results(input_message.content, image, process_time, save) @@ -8,7 +8,6 @@ from typing import Any import filetype import requests from django.core.files import File -from replicate.exceptions import ModelError from messages.models import Message from ml_model.exceptions import FileNotProvided, RequestBlocked, GenerationException @@ -69,12 +68,7 @@ class Runway(SimpleService): 'image': media } start_time = time.time() - try: - video = replicate_run(f'runwayml/{version}', callback_data) - except ModelError as exc: - if any(error in str(exc) for error in ('E005', 'E006', 'sexual')): - raise RequestBlocked - raise GenerationException from exc + video = replicate_run(f'runwayml/{version}', callback_data) process_time = timedelta(seconds=(time.time() - start_time)) self.handle_invoice(input_message.content_object.model, duration=duration, version=version) msgs = self.save_results(input_message.content, process_time, video, save) @@ -8,7 +8,6 @@ from typing import Any import filetype import requests from django.core.files import File -from replicate.exceptions import ModelError from messages.models import Message from ml_model.exceptions import RequestBlocked, GenerationException, FileExtensionNotSupported @@ -94,14 +93,9 @@ class Seedance(SimpleService): ) callback_data.update({f'reference_{reference_type}': [file.url]}) start_time = time.time() - try: - video = replicate_run( - f'bytedance/{version}', callback_data - ) - except ModelError as exc: - if any(error in str(exc) for error in ('E005', 'E006', 'sexual')): - raise RequestBlocked - raise GenerationException from exc + video = replicate_run( + f'bytedance/{version}', callback_data + ) process_time = timedelta(seconds=(time.time() - start_time)) self.handle_invoice(input_message.content_object.model, resolution=resolution, duration=duration, version=version, generation_type=generation_type) msgs = self.save_results(input_message.content, process_time, video, save) @@ -6,7 +6,6 @@ from typing import Any import requests from django.core.files import File -from replicate.exceptions import ModelError from messages.models import Message from ml_model.exceptions import RequestBlocked, GenerationException @@ -52,12 +51,7 @@ class Seedream(SimpleService): } if image := input_message.file: callback_data.update({'image_input': [image.url]}) - try: - images = replicate_run('bytedance/seedream-5-lite', callback_data) - except ModelError as exc: - if any(error in str(exc) for error in ('E005', 'E006', 'sexual')): - raise RequestBlocked - raise GenerationException from exc + images = replicate_run('bytedance/seedream-5-lite', callback_data) process_time = timedelta(seconds=(time.time() - start_time)) self.handle_invoice(input_message.content_object.model) msgs = self.save_results(input_message.content, images, process_time, save) @@ -8,7 +8,6 @@ from typing import Any import filetype import requests from django.core.files import File -from replicate.exceptions import ModelError from messages.models import Message from ml_model.exceptions import RequestBlocked, GenerationException @@ -56,12 +55,7 @@ class Veo(SimpleService): input_message.file.close() callback_data.update({'image': image}) start_time = time.time() - try: - video = replicate_run(f'google/{version}', callback_data) - except ModelError as exc: - if any(error in str(exc) for error in ('E005', 'E006', 'sexual')): - raise RequestBlocked - raise GenerationException from exc + video = replicate_run(f'google/{version}', callback_data) process_time = timedelta(seconds=(time.time() - start_time)) self.handle_invoice(input_message.content_object.model, version) msgs = self.save_results(input_message.content, process_time, video, save) @@ -8,7 +8,6 @@ from typing import Any import filetype import requests from django.core.files import File -from replicate.exceptions import ModelError from messages.models import Message from ml_model.exceptions import FileExtensionNotSupported, CorruptedFileError, GenerationException @@ -64,12 +63,7 @@ class Wan_Lite(SimpleService): input_message.file.close() callback_data.update({'image': image}) start_time = time.time() - try: - video = replicate_run('wan-video/wan-2.2-5b-fast', callback_data) - except ModelError as exc: - if 'image file' in exc.prediction.error: - raise CorruptedFileError from exc - raise GenerationException from exc + video = replicate_run('wan-video/wan-2.2-5b-fast', callback_data) process_time = timedelta(seconds=(time.time() - start_time)) self.handle_invoice(input_message.content_object.model, resolution=resolution) msgs = self.save_results(input_message.content, process_time, video, save) @@ -15,10 +15,25 @@ import replicate import requests from celery import shared_task from deepl.translator import TextResult +from django.utils.translation import gettext as _ +from replicate.exceptions import ModelError from requests import Response from backend import settings -from ml_model.exceptions import DeploymentDisabled, ModelTimeoutError, GenerationException, RequestBlocked +from ml_model.exceptions import ( + CorruptedFileError, + DeploymentDisabled, + ExceededContextLengthError, + FaceNotFoundError, + GenerationException, + ImageAnalysisError, + ImageContentNotFound, + InvalidParameterError, + ModelCouldNotInterpretPrompt, + ModelTimeoutError, + PredictionInterruptedError, + RequestBlocked, +) from ml_model.utils import count_openrouter_tokens from poller.models import Proxy @@ -95,10 +110,35 @@ def transcript_audio(payload: dict[str, Any]): @shared_task def replicate_run(callback_url: str, payload: dict[str, Any]): replicate_client = replicate.Client(settings.REPLICATE_API_KEY) - return replicate_client.run( - ref=callback_url, - input=payload, - ) + try: + return replicate_client.run( + ref=callback_url, + input=payload, + ) + except ModelError as exc: + prediction_error = getattr(getattr(exc, 'prediction', None), 'error', '') or '' + error_text = str(exc) + if any(error in error_text for error in ('E005', 'E006', 'sexual', 'NSFW')): + raise RequestBlocked + if 'PA' in error_text: + raise PredictionInterruptedError + if prediction_error == 'No image content found in response': + raise ImageContentNotFound + if prediction_error == 'Failed to generate image.': + raise ModelCouldNotInterpretPrompt from exc + if prediction_error == '400': + raise ModelCouldNotInterpretPrompt from exc + if prediction_error == 'facexlib align face fail': + raise FaceNotFoundError + if prediction_error and 'image file' in prediction_error: + raise CorruptedFileError from exc + if 'lyrics is too long' in error_text: + raise InvalidParameterError(_('Lyrics is too long')) + if 'INPUT_ANALYSIS_FAILURE' in error_text: + raise ImageAnalysisError + if 'PROMPT_TOO_LONG' in error_text: + raise ExceededContextLengthError + raise GenerationException from exc @shared_task