@@ -32,4 +32,5 @@ from ml_model.services.sdxlemoji import Sdxlemoji from ml_model.services.stablediffusion import Stablediffusion from ml_model.services.upscaleai import Upscaleai from ml_model.services.vicuna import Vicuna +from ml_model.services.wan import Wan from ml_model.services.whisper import Whisper @@ -0,0 +1,47 @@ +import time +from datetime import timedelta +from decimal import Decimal +from io import BytesIO + +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 + + +class Wan(SimpleService): + """ + Wan Service + contains abstract method make, which makes a generation + """ + + TOKENS_COST = { + '480p': Decimal('25'), + '720p': Decimal('50') + } + + def calculate_price(self, resolution: str) -> Decimal: + return self.TOKENS_COST[resolution].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', '720p') + callback_data = dict({'prompt': input_message.content, 'resolution': resolution, **input_message.info}) + start_time = time.time() + video = replicate_run('wan-video/wan-2.2-t2v-fast', callback_data) + process_time = timedelta(seconds=(time.time() - start_time)) + self.handle_invoice(input_message.content_object.model, resolution) + msgs = self.save_results(input_message.content, process_time, video, save) + return msgs