@@ -35,6 +35,7 @@ from ml_model.services.lightning import Lightning from ml_model.services.llama import Llama from ml_model.services.logoai import Logoai from ml_model.services.lyria import Lyria +from ml_model.services.ltx import Ltx from ml_model.services.midjourney import Midjourney from ml_model.services.minimaxvideo import Minimaxvideo from ml_model.services.minimaxmusic import Minimaxmusic @@ -0,0 +1,79 @@ +import time +from datetime import timedelta +from decimal import Decimal +from io import BytesIO +from typing import Any + +import requests +from django.core.files import File + +from messages.models import Message +from ml_model.services.base import SimpleService +from ml_model.tasks import replicate_run + +from payments.exceptions.insufficient_balance import InsufficientBalance +from payments.selectors.payment_plan_selector import PaymentPlanSelector + + +class Ltx(SimpleService): + TOKENS_COST = { + '1080p': Decimal('12'), + '2k': Decimal('24'), + '4k': Decimal('48'), + } + + @classmethod + def predict_price(cls, content: str, file_exists: bool, info: dict[str, Any]) -> Decimal | None: + resolution = info['resolution'] + duration = info['duration'] + return (cls.TOKENS_COST[resolution] * duration).quantize(Decimal('0.1'), rounding='ROUND_UP') + + def calculate_price(self, resolution: str, duration: int) -> Decimal: + return (self.TOKENS_COST[resolution] * duration).quantize(Decimal('0.1'), rounding='ROUND_UP') + + def save_results(self, content: str, t: timedelta, video: str, save: bool = True) -> list[Message]: + msg = Message( + content=content, + content_object=self.store, + elapsed_time=t, + file=File(BytesIO(requests.get(video).content), '.mp4'), + ) + if save: + return Message.objects.bulk_create([msg]) + return [msg] + + def make(self, input_message: Message, save: bool = True) -> list[Message]: + resolution = input_message.info.pop('resolution', '1080p') + duration = input_message.info.pop('duration', 6) + if (balance := PaymentPlanSelector(self.store.user).get_current_balance()) < ( + cost := self.TOKENS_COST[resolution] * duration + ): + raise InsufficientBalance(balance, cost) + camera_motion = { + 'Без движения камеры': 'none', + 'Приближение камеры': 'dolly_in', + 'Удаление камеры': 'dolly_out', + 'Движение камеры влево': 'dolly_left', + 'Движение камеры вправо': 'dolly_right', + 'Подъём камеры': 'jib_up', + 'Опускание камеры': 'jib_down', + 'Статичная камера': 'static', + 'Смена фокуса': 'focus_shift', + } + callback_data = dict( + { + 'prompt': input_message.content, + 'resolution': resolution, + 'duration': duration, + 'camera_motion': camera_motion[input_message.info.pop('camera_motion', 'Без движения камеры')], + **input_message.info, + } + ) + if input_message.file: + callback_data.update({'image': input_message.file.url}) + start_time = time.time() + video = replicate_run('lightricks/ltx-2.3-fast', callback_data) + process_time = timedelta(seconds=(time.time() - start_time)) + self.handle_invoice(input_message.content_object.model, resolution=resolution, duration=duration) + msgs = self.save_results(input_message.content, process_time, video, save) + return msgs