@@ -6,6 +6,7 @@ from ml_model.services.deepl import Deepl from ml_model.services.epicphotogasm import Epicphotogasm from ml_model.services.djourney import Djourney from ml_model.services.flux import Flux +from ml_model.services.iconic import Iconic from ml_model.services.kandinsky import Kandinsky from ml_model.services.llama import Llama from ml_model.services.logoai import Logoai @@ -0,0 +1,156 @@ +import time + +import requests +from datetime import timedelta +from decimal import Decimal +from io import BytesIO + +from messages.models import Message +from ml_model.exceptions.external_api import ExternalAPIException +from ml_model.models import ModelCategory, ModelInput, ModelParameter +from ml_model.services.base import SimpleService +from ml_model.tasks import replicate_run + +from django.core.files import File + + +class Iconic(SimpleService): + """ + Iconic Service + contains abstract method make, which makes a generation + """ + title = 'Iconic' + description = 'Нейросеть, способная генерировать картинки из вашего текста' + category = ModelCategory(title='Изображения', slug='images') + versions = [] + inputs = [ + ModelInput(type=ModelInput.TypeChoices.TEXT, required=True), + ModelInput(type=ModelInput.TypeChoices.IMAGE) + ] + parameters = [ + ModelParameter( + name='Модель', + key='model', + type=ModelParameter.TypeChoices.LIST, + values={ + 'availables': [ + 'dev', + 'schnell', + ], + 'default': 'dev', + }, + ), + ModelParameter( + name='Ширина', + key='width', + type=ModelParameter.TypeChoices.INTRANGE, + values={'start': 256, 'end': 1440, 'step': 128, 'default': 1024}, + ), + ModelParameter( + name='Высота', + key='height', + type=ModelParameter.TypeChoices.INTRANGE, + values={'start': 256, 'end': 1440, 'step': 128, 'default': 1024}, + ), + ModelParameter( + name='Количество изображений', + key='num_outputs', + type=ModelParameter.TypeChoices.INTRANGE, + values={'start': 1, 'end': 4, 'step': 1, 'default': 1}, + ), + ModelParameter( + name='Соотношение сторон', + key='aspect_ratio', + type=ModelParameter.TypeChoices.LIST, + values={ + 'availables': [ + '1:1', + '16:9', + '21:9', + '3:2', + '2:3', + '4:5', + '5:4', + '3:4', + '4:3', + '9:16', + '9:21', + 'custom' + ], + 'default': '1:1', + }, + ), + ModelParameter( + name='Точность запроса', + key='guidance_scale', + type=ModelParameter.TypeChoices.FLOATRANGE, + values={'start': 0.0, 'end': 10.0, 'step': 1.0, 'default': 3.5}, + ), + ModelParameter( + name='Качество вывода', + key='output_quality', + type=ModelParameter.TypeChoices.INTRANGE, + values={'start': 0, 'end': 100, 'step': 1, 'default': 90}, + ), + ModelParameter( + name='Шаги предобработки', + key='num_inference_steps', + type=ModelParameter.TypeChoices.INTRANGE, + values={'start': 1, 'end': 50, 'step': 1, 'default': 28}, + ), + ModelParameter( + name='Отключить проверку безопасности', + key='disable_safety_checker', + type=ModelParameter.TypeChoices.BOOL, + values={'default': False}, + ), + ] + + PRICE = Decimal('0.980') + + _CALLBACK = ( + 'miike-ai/flux-ico' + ':478cae37f1aec0fde7977fdd54b272aaeabede7d8060801841920c16306369a9' + ) + + def calculate_price(self, process_time: timedelta) -> Decimal: + total_seconds = Decimal(process_time.total_seconds()) + return self.PRICE * total_seconds + + def save_results( + self, prompt: str, images: list, time: timedelta, save: bool = True + ) -> list[Message]: + messages: list[Message] = [] + for image in images: + messages.append( + Message( + content_object=self.store, + elapsed_time=time, + content=prompt, + file=File(BytesIO(requests.get(image).content), '.png') + ) + ) + if save: + return Message.objects.bulk_create(messages) + return messages + + def make(self, input_message: Message, save: bool = True) -> list[Message]: + start_time = time.time() + translated_prompt = self.translate_prompt(input_message.content) + callback_data = dict( + { + 'prompt': f'ICO, Create a flat icon for {translated_prompt}', + **input_message.info, + } + ) + 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() + 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) + return msgs