@@ -12,6 +12,7 @@ from ml_model.services.logoai import Logoai from ml_model.services.mistral import Mistral from ml_model.services.musicgen import Musicgen from ml_model.services.openjourney import Openjourney +from ml_model.services.pulid import Pulid from ml_model.services.sdxlemoji import Sdxlemoji from ml_model.services.stablediffusion import Stablediffusion from ml_model.services.upscaleai import Upscaleai @@ -0,0 +1,143 @@ +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 Pulid(SimpleService): + """ + PuLID Service + contains abstract method make, which makes a generation + """ + title = 'PuLID' + description = 'Нейросеть, способная генерировать фотографии из вашего текста' + category = ModelCategory(title='Изображения', slug='images') + versions = [] + inputs = [ + ModelInput(type=ModelInput.TypeChoices.TEXT), + ModelInput(type=ModelInput.TypeChoices.IMAGE, required=True) + ] + parameters = [ + ModelParameter( + name='Негативный промпт', + key='negative_prompt', + type=ModelParameter.TypeChoices.STR, + ), + ModelParameter( + name='CFG Scale', + key='cfg_scale', + type=ModelParameter.TypeChoices.FLOATRANGE, + values={'start': 1.0, 'end': 1.5, 'step': 0.1, 'default': 1.2}, + ), + ModelParameter( + name='Шаги обработки', + key='num_steps', + type=ModelParameter.TypeChoices.INTRANGE, + values={'start': 1, 'end': 100, 'step': 1, 'default': 4}, + ), + ModelParameter( + name='Ширина', + key='image_width', + type=ModelParameter.TypeChoices.INTRANGE, + values={'start': 1024, 'end': 2024, 'step': 128, 'default': 1024}, + ), + ModelParameter( + name='Высота', + key='image_height', + type=ModelParameter.TypeChoices.INTRANGE, + values={'start': 1024, 'end': 2024, 'step': 128, 'default': 1024}, + ), + ModelParameter( + name='Количество изображений', + key='num_samples', + type=ModelParameter.TypeChoices.INTRANGE, + values={'start': 1, 'end': 8, 'step': 1, 'default': 4}, + ), + ModelParameter( + name='Шкала идентичности', + key='identity_scale', + type=ModelParameter.TypeChoices.FLOATRANGE, + values={'start': 0, 'end': 5.0, 'step': 0.1, 'default': 0.8}, + ), + ModelParameter( + name='Качество вывода', + key='output_quality', + type=ModelParameter.TypeChoices.INTRANGE, + values={'start': 0, 'end': 100, 'step': 1, 'default': 80}, + ), + ModelParameter( + name='Режим генерации', + key='generation_mode', + type=ModelParameter.TypeChoices.LIST, + values={ + 'availables': [ + 'fidelity', + 'extremely style', + ], + 'default': 'fidelity', + }, + ), + ] + + PRICE = Decimal('0.340') + + _CALLBACK = ( + 'zsxkib/pulid' + ':43d309c37ab4e62361e5e29b8e9e867fb2dcbcec77ae91206a8d95ac5dd451a0' + ) + + def calculate_price(self, process_time: timedelta) -> Decimal: + return self.PRICE * Decimal(process_time.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() + callback_data = dict( + { + 'prompt': f'portrait, {self.translate_prompt(input_message.content)}', + 'output_format': 'png', + 'negative_prompt': 'flaws in the eyes, flaws in the face, flaws, lowres, non-HDRi, low quality, ' + 'worst quality, artifacts noise, text, watermark, glitch, deformed, mutated, ' + 'ugly, disfigured, hands, low resolution, partially rendered objects, deformed ' + 'or partially rendered eyes, deformed eyeballs, cross-eyed, blurry, ' + f'{input_message.info.pop('negative_prompt', '')}', + **input_message.info, + } + ) + 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() + 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