@@ -749,8 +749,6 @@ msgid "Available only in paid plan" msgstr "Доступно только в платном тарифе" #: ml_model/exceptions.py:163 -#, fuzzy -#| msgid "Image analysis error. Please try another image." msgid "Face not found in the image. Please try another image with a face." msgstr "Не найдено лицо на картинке. Попробуйте другую картинку с лицом." @@ -35,3 +35,5 @@ def calculate_predict_price(request, body: PredictPriceInputSchema): return PredictPriceSchema(price=predicted_price) + + @@ -107,7 +107,8 @@ class Chatgpt_5_4(Chatgpt): model_name = info.pop('version', 'gpt-5.4') user_system_prompt = info.pop('system_prompt', '') plan_info = self.store.user.payment_plan - is_free_plan = plan_info and plan_info.plan.price <= 0 + is_regular_user = self.store.user.account_type == 'regular' + is_free_plan = is_regular_user and plan_info and plan_info.plan.price <= 0 if is_free_plan and model_name == 'gpt-5.4-pro': raise PaidPlanRequiredError() if is_free_plan: @@ -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) @@ -12,24 +12,32 @@ from tools.public_api.models import APIStore class Deepseek(SimpleService): TOKENS_COST = { - 'deepseek/deepseek-chat': { - 'input': Decimal('390') / 1_000_000, - 'output': Decimal('390') / 1_000_000, + # 'deepseek/deepseek-chat': { + # 'input': Decimal('390') / 1_000_000, + # 'output': Decimal('390') / 1_000_000, + # }, + # 'deepseek/deepseek-r1': { + # 'input': Decimal('900') / 1_000_000, + # 'output': Decimal('900') / 1_000_000, + # }, + # 'deepseek/deepseek-r1:free': { + # 'input': Decimal('0'), + # 'output': Decimal('0'), + # }, + 'deepseek/deepseek-v4-pro': { + 'input': Decimal('1050') / 1_000_000, + 'output': Decimal('2200') / 1_000_000, }, - 'deepseek/deepseek-r1': { - 'input': Decimal('900') / 1_000_000, - 'output': Decimal('900') / 1_000_000, - }, - 'deepseek/deepseek-r1:free': { - 'input': Decimal('0'), - 'output': Decimal('0') + 'deepseek/deepseek-v4-flash': { + 'input': Decimal('100') / 1_000_000, + 'output': Decimal('175') / 1_000_000, }, } - PRICE_BIAS = Decimal('0.05') def calculate_price(self, version: str, input_tokens: int, output_tokens: int) -> Decimal: price_map = self.TOKENS_COST[version] - price = input_tokens * price_map['input'] + output_tokens * price_map['output'] + self.PRICE_BIAS + price = input_tokens * price_map['input'] + output_tokens * price_map['output'] + return price.quantize(Decimal('0.1'), rounding='ROUND_UP') def save_results(self, content: str, t: timedelta, save: bool = True) -> list[Message]: @@ -47,15 +55,17 @@ class Deepseek(SimpleService): def make(self, input_message: Message, save: bool = True) -> list[Message]: info = input_message.info.copy() version = info.pop('version') - system_prompt = input_message.info.pop('system_prompt', '') + system_prompt = info.pop('system_prompt', '') + callback_data = { - **input_message.info, + **info, } messages = [ {'role': 'system', 'content': system_prompt}, *self.get_chat_history(), {'role': 'user', 'content': input_message.content} ] + start_time = time.time() result = openrouter_run(version, messages, callback_data, 'Deepseek') process_time = timedelta(seconds=(time.time() - start_time)) @@ -66,8 +76,11 @@ class Deepseek(SimpleService): output_tokens=result[2], ) msgs = self.save_results(result[0], process_time) + return msgs + + def get_chat_history(self, message_limit: int = 10, max_character_limit: int = 1500) -> list[dict[str, str | list]]: if isinstance(self.store, Chat): air_messages = list( @@ -100,3 +113,4 @@ class Deepseek(SimpleService): while character_length > max_character_limit: character_length -= len(memory.pop(0)['content']) return memory + @@ -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) @@ -8,7 +8,6 @@ import requests from django.core.files import File from messages.models import Message -from ml_model.exceptions import RequestBlocked, PredictionInterruptedError, GenerationException from ml_model.models import ( NeuronModel, ) @@ -87,17 +86,10 @@ class Flux(SimpleService): **input_message.info, } ) - try: - runner = replicate_run( - f'{self._CALLBACK_BASE}{callback_data.get("version", "flux-schnell")}', - callback_data, - ) - except Exception as exc: - if any(error in str(exc) for error in ('E005', 'E006', 'sexual', 'NSFW')): - raise RequestBlocked - elif 'PA' in str(exc): - raise PredictionInterruptedError - raise GenerationException from exc + runner = replicate_run( + f'{self._CALLBACK_BASE}{callback_data.get("version", "flux-schnell")}', + 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) @@ -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 @@ -4,10 +4,16 @@ from datetime import timedelta from decimal import Decimal from io import BytesIO from typing import Any, Dict, Iterator +from pathlib import Path import filetype from django.db.models.fields.files import FieldFile -from PIL import Image +from PIL import Image, DecompressionBombError +from ml_model.services.FileService import FileProcessingService +from ml_model.exceptions import FileExtensionNotSupported, CorruptedFileError, ImageTooLargeError +from ml_model.services.EmbeddingService import EmbeddingService + +from poller.models import Proxy from messages.models import Message from ml_model.services.base import SimpleService @@ -35,8 +41,14 @@ class Grok(SimpleService): }, # 1M tokens } + TOOLS_TOKEN_COSTS = {'text-embedding-3-small': {'output': Decimal('0.00001')}} + + SUPPORTED_EXTENSIONS = ['PDF', 'DOC', 'DOCX', 'XLSX', 'JPG', 'JPEG', 'PNG', 'WEBP'] + + MAX_PIXELS = 178956970 + def calculate_price( - self, version: str, input_tokens: int, output_tokens: int, image: FieldFile + self, version: str, input_tokens: int, output_tokens: int, image: FieldFile, embedding_tokens: int ) -> Decimal: price_map = self.TOKENS_COST[version.split('/')[1]] price = ( @@ -44,6 +56,8 @@ class Grok(SimpleService): ) if image: price += price_map['input_imgs'] / 1_000 + if embedding_tokens: + price += embedding_tokens * self.TOOLS_TOKEN_COSTS['text-embedding-3-small']['output'] return price.quantize(Decimal('0.1'), rounding='ROUND_UP') def save_results(self, content: Iterator[Any], t: timedelta, save: bool = True) -> list[Message]: @@ -64,31 +78,75 @@ class Grok(SimpleService): callback_data = {**input_message.info} messages = self.get_chat_history() messages.append({'role': 'user', 'content': input_message.content}) - image = input_message.file - if image: - kind = filetype.guess(image.read(20)) - mime = kind.mime if kind else 'application/octet-stream' - normalized_image = Image.open(image) - format = 'jpeg' if kind.extension == 'jpg' else kind.extension - buf = BytesIO() - normalized_image.save(buf, format=format) - image_url = f'data:{mime};base64,{base64.b64encode(buf.getvalue()).decode("utf-8")}' - buf.close() - messages[-1]['content'] = [ - {'type': 'text', 'text': input_message.content}, - {'type': 'image_url', 'image_url': {'url': image_url}}, - ] - result = openrouter_run(version, messages, callback_data, 'Grok') - process_time = timedelta(seconds=(time.time() - start_time)) - self.handle_invoice( - input_message.content_object.model, - version=version, - input_tokens=result[1], - output_tokens=result[2], - image=image, - ) - msgs = self.save_results(result[0], process_time) - return msgs + embedding_tokens = 0 + if input_message.file: + file_service = FileProcessingService + file_bytes = input_message.file.read() + kind = filetype.guess(file_bytes[:550]) + image = input_message.file + if not kind: + if Path(input_message.file.name).suffix[1:].upper() not in self.SUPPORTED_EXTENSIONS: + raise FileExtensionNotSupported(self.SUPPORTED_EXTENSIONS) + raise CorruptedFileError + + raw_file_extension = kind.extension + file_extension = file_service.get_file_extension(raw_file_extension, file_bytes) + if file_extension in ('pdf', 'doc', 'docx', 'xlsx'): + text = file_service.get_file_data(file_extension, file_bytes) + chunks = EmbeddingService.split_text_to_chunks(text) + + if len(text) > 20_000: + for proxy in Proxy.objects.all(): + document_name = chunks[0].partition(f':{chr(10)}')[2].split(f'{chr(10)}')[0][:100] + embedding_tokens, file_data = EmbeddingService.get_large_file_data( + self.store.messages.first().pk, + chunks, + proxy, + input_message.content, + model='text-embedding-3-small', + index_name='ml_model-index-1536', + ) + messages[-1]['content'] = EmbeddingService.make_embeddings_prompt( + document_name=document_name, + section_texts=file_data, + question=input_message.content, + ) + else: + messages[-1]['content'] = ( + f'Используй системный промпт. Содержание файла: ' + f'{chunks}. Вопрос: {input_message.content}' + ) + + + elif file_extension in ('jpg', 'jpeg', 'png', 'webp'): + mime = kind.mime if kind else 'application/octet-stream' + normalized_image = Image.open(input_message.file) + format = 'jpeg' if kind.extension == 'jpg' else kind.extension + buf = BytesIO() + normalized_image.save(buf, format=format) + image_url = f'data:{mime};base64,{base64.b64encode(buf.getvalue()).decode("utf-8")}' + buf.close() + messages[-1]['content'] = [ + {'type': 'text', 'text': input_message.content}, + {'type': 'image_url', 'image_url': {'url': image_url}}, + ] + else: + raise FileExtensionNotSupported(self.SUPPORTED_EXTENSIONS) + + result = openrouter_run(version, messages, callback_data, 'Grok') + process_time = timedelta(seconds=(time.time() - start_time)) + self.handle_invoice( + input_message.content_object.model, + version=version, + input_tokens=result[1], + output_tokens=result[2], + image=image, + ) + msgs = self.save_results(result[0], process_time) + return msgs + + + def get_chat_history(self, message_limit: int = 10, max_character_limit: int = 1500) -> list[dict[str, str | list]]: if isinstance(self.store, Chat): @@ -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) @@ -7,7 +7,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 @@ -93,14 +92,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) @@ -39,7 +39,6 @@ class UnsupportedSize(Exception): self.current_size | self.required_size ) - class ModelTimeoutError(Exception): def __str__(self): return _('The model is not responding') @@ -1,5 +1,4 @@ from typing import List, Optional, Any - from ninja import ModelSchema, Schema from pydantic import condecimal @@ -15,11 +15,26 @@ 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.adapters.bytedance_model_ark import BytedanceContentType, BytedanceModelArkAdapter -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 @@ -96,10 +111,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 @@ -5851,13 +5851,13 @@ files = [ [[package]] name = "yookassa" -version = "2.5.0" +version = "3.10.1" description = "YooKassa API SDK Python Library" optional = false python-versions = "*" groups = ["main"] files = [ - {file = "yookassa-2.5.0.tar.gz", hash = "sha256:5ddb279d6e867c74b66549e3096196606b5f04bf4927bde2513072b7a08ee3ff"}, + {file = "yookassa-3.10.1.tar.gz", hash = "sha256:6eddff428be5eb86c001c79c334ce8ad265b52b315baba70dc5f19bec0f25955"}, ] [package.dependencies] @@ -5945,4 +5945,4 @@ testing = ["coverage[toml]", "zope.event", "zope.testing"] [metadata] lock-version = "2.1" python-versions = "^3.12" -content-hash = "062dda7a1043076481446c18a95fa53ce5e21da80cdbe045621204ecf8a72811" +content-hash = "6d8d326ffa25ad4bc05f2319238735e45ac25671682a40765e79d603685635d4" @@ -13,7 +13,7 @@ django = "5.0.*" django-cors-headers = "^4.2.0" djangorestframework-simplejwt = "^5.2.2" djangorestframework = "^3.14.0" -yookassa = "^2.4.0" +yookassa = "^3.10.1" environs = "^9.5.0" dj-rest-auth = "^4.0.1" django-celery-beat = "^2.5.0"