@@ -1116,8 +1116,8 @@ msgid "1080p output is not supported for Seedance Dreamina 2.0 Fast." msgstr "1080р разрешение не поддерживается для Seedance Dreamina 2.0 Fast." #: ml_model/services/seedream.py:90 -msgid "3K output is not supported for Seedream 4.5" -msgstr "3К разрешение не поддерживается для Seedream 4.5" +msgid "3K output is not supported for this model" +msgstr "3К разрешение не поддерживается для этой модели" #: ml_model/services/upscaleai.py:124 msgid "No image given for improving" @@ -21,20 +21,20 @@ from payments.selectors.payment_plan_selector import PaymentPlanSelector class Seedream(SimpleService): TOKEN_COST = { - "seedream-boosted": { - "2K": Decimal('25'), - "3K": Decimal('50'), - "4K": Decimal('100'), + 'seedream-boosted': { + '2K': Decimal('25'), + '3K': Decimal('50'), + '4K': Decimal('100'), }, - "seedream-4.5": { - "2K": Decimal('25'), - "4K": Decimal('100'), + 'seedream-4.5': { + '2K': Decimal('25'), + '4K': Decimal('100'), }, } VERSION_MAPPING = { - "seedream-boosted": "seedream-5-0-260128", - "seedream-4.5": "seedream-4-5-251128", + 'seedream-boosted': 'seedream-5-0-260128', + 'seedream-4.5': 'seedream-4-5-251128', } OPTIMIZATION_PROMPT = """ @@ -74,23 +74,22 @@ class Seedream(SimpleService): def make(self, input_message: Message, save: bool = True) -> list[Message]: start_time = time.time() - if (balance := PaymentPlanSelector(self.store.user).get_current_balance()) < (predicted := self.predict_price( - input_message.content, - input_message.file is not None, - input_message.info, - )): - raise InsufficientBalance(balance, predicted) - version = input_message.info.pop('version', None) + size = input_message.info.get('size', '2K') + if version is None or version not in self.TOKEN_COST: raise ModelVersionNotAvailable(version, self.TOKEN_COST) - size = input_message.info.get('size', '2K') if version == 'seedream-4.5' and size == '3K': - raise InvalidParameterError(_('3K output is not supported for Seedream 4.5')) + raise InvalidParameterError(_('3K output is not supported for this model')) + + if (balance := PaymentPlanSelector(self.store.user).get_current_balance()) < ( + predicted := self.calculate_price(version, size) + ): + raise InsufficientBalance(balance, predicted) callback_data = { - 'prompt': f"{input_message.content}\n{self.OPTIMIZATION_PROMPT}", + 'prompt': f'{input_message.content}\n{self.OPTIMIZATION_PROMPT}', 'watermark': False, **input_message.info, } @@ -1,120 +1,38 @@ -import logging -import time -import uuid -from datetime import timedelta -from decimal import Decimal -from io import BytesIO from typing import Any -import httpx -from django.conf import settings -from django.core.files import File - from messages.models import Message -from ml_model.exceptions import ModelVersionNotAvailable, RequestBlocked -from ml_model.services.base import SimpleService -from poller.models import Proxy - -logger = logging.getLogger(__name__) +from ml_model.exceptions import ModelVersionNotAvailable +from ml_model.services.seedream import Seedream -class Stablediffusion(SimpleService): +class Stablediffusion(Seedream): """ Stablediffusion Service contains abstract method make, which makes a generation """ - MODELS = ['sd3', 'sd3-turbo', 'sd3-medium'] - MODELS_LINKS = { - 'sd3': 'stable-diffusion-3.5-large', - 'sd3-turbo': 'stable-diffusion-3.5-large-turbo', - 'sd3-medium': 'stable-diffusion-3.5-medium', + PROXY_VERSION_MAPPING = { + 'sd3': 'seedream-boosted', + 'sd3-medium': 'seedream-4.5', + 'sd3-turbo': 'seedream-4.5', } - def calculate_price(self, input_message: Message) -> Decimal: - if input_message.info.get('version') == 'sd3': - return Decimal('32.5') - elif input_message.info.get('version') == 'sd3-turbo': - return Decimal('20') - elif input_message.info.get('version') == 'sd3-medium': - return Decimal('17.5') - @classmethod - def predict_price(cls, content: str, file_exists: bool, info: dict[str, Any]) -> Decimal | None: - version = info.get('version') - if version == 'sd3': - return Decimal('32.5') - elif version == 'sd3-turbo': - return Decimal('20') - elif version == 'sd3-medium': - return Decimal('17.5') - return None + def _remap_info(cls, info: dict[str, Any]) -> dict[str, Any]: + new_info = info.copy() + version = info.get('version', 'sd3') + if version not in cls.PROXY_VERSION_MAPPING: + raise ModelVersionNotAvailable(version, cls.PROXY_VERSION_MAPPING) + new_info['version'] = cls.PROXY_VERSION_MAPPING[version] + return new_info - def save_results( - self, - input_prompt: str, - link: str, - t: timedelta, - save: bool = True, - ) -> list[Message]: - out: list[Message] = [] - out.append( - Message( - content_object=self.store, - elapsed_time=t, - content=input_prompt, - file=File( - BytesIO(httpx.get(link).content), - f'{uuid.uuid4()}.png', - ), - ) - ) - if save: - return Message.objects.bulk_create(out) - return out + @classmethod + def predict_price(cls, content: str, file_exists: bool, info: dict[str, Any]): + try: + return super().predict_price(content, file_exists, cls._remap_info(info)) + except ModelVersionNotAvailable: + return None def make(self, input_message: Message, save: bool = True) -> list[Message]: - start_time = time.time() - info = input_message.info.copy() - version = info.get('version') - if version is None or version not in self.MODELS: - raise ModelVersionNotAvailable(version, self.MODELS) - model_name = self.MODELS_LINKS[version] - translated_prompt = self.translate_prompt(input_message.content) - callback_data = { - 'prompt': translated_prompt, - 'aspect_ratio': input_message.info.get('aspect_ratio', '1:1'), - 'output_quality': 100, - 'output_format': 'png', - } - link = None - for proxy in Proxy.objects.all(): - with httpx.Client( - headers={ - 'Authorization': f'Bearer {settings.REPLICATE_API_KEY}', - 'Prefer': 'wait', - 'Content-Type': 'application/json', - }, - timeout=600, - proxy=f'{proxy.protocol}://{proxy.address}', - ) as client: - result = client.post( - f'https://api.replicate.com/v1/models/stability-ai/{model_name}/predictions', - json={'input': callback_data}, - ).json() - while result['status'] not in ('succeeded', 'failed', 'canceled'): - result = client.get(result['urls']['get']).json() - if result['status'] in ('failed', 'canceled'): - if 'E005' in result['logs']: - raise RequestBlocked - logger.error(result['logs']) - raise Exception('No answer from Stable Diffusion, please retry later') - link = result['output'] - - process_time = timedelta(seconds=(time.time() - start_time)) - - if not link: - raise Exception('No answer from Stable Diffusion, please retry later') - self.handle_invoice(input_message.content_object.model, input_message) - msgs = self.save_results(input_message.content, link, process_time, save) - return msgs + input_message.info = self._remap_info(input_message.info) + return super().make(input_message, save) \ No newline at end of file