@@ -48,3 +48,4 @@ from ml_model.services.minimaxvideo import Minimaxvideo from ml_model.services.kling import Kling from ml_model.services.fluxkrea import Fluxkrea from ml_model.services.ideogram import Ideogram +from ml_model.services.leonardo import Leonardo @@ -0,0 +1,82 @@ +import base64 +import time +from datetime import timedelta +from decimal import Decimal +from io import BytesIO + +import filetype +import requests +from django.core.files import File + +from messages.models import Message +from ml_model.models import ( + NeuronModel, +) +from ml_model.services.base import SimpleService +from ml_model.tasks import replicate_run + + +class Leonardo(SimpleService): + """ + Flux Service + contains abstract method make, which makes a generation + """ + + TOKENS_COST = { + 'lucid-origin': { + 'input_imgs': Decimal('1.5'), + }, # 1k images + } + + _CALLBACK_BASE = 'leonardoai/' + + def calculate_price(self, input_message: Message, version: str) -> Decimal: + price_map = self.TOKENS_COST[version] + price = price_map['input_imgs'] / 1_000 + if num_images := input_message.info.get('num_images'): + price = price * num_images + return price.quantize(Decimal('0.1'), rounding='ROUND_UP') + + @property + def neuron_model(self): + return NeuronModel.objects.get(title='Flux') + + 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() + version = input_message.info.get('version') + callback_data = dict( + { + 'prompt': self.translate_prompt(input_message.content), + **input_message.info, + } + ) + 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(input_message.content_object.model, input_message=input_message, version=version) + msgs = self.save_results(input_message.content, images, process_time, save) + return msgs