@@ -7,6 +7,7 @@ from ml_model.services.deepl import Deepl from ml_model.services.deepseek import Deepseek from ml_model.services.djourney import Djourney from ml_model.services.epicphotogasm import Epicphotogasm +from ml_model.services.elevenlabs import Elevenlabs from ml_model.services.flux import Flux from ml_model.services.flux_2 import Flux_2 from ml_model.services.fluxkrea import Fluxkrea @@ -41,6 +42,7 @@ from ml_model.services.nanobanana import Nanobanana from ml_model.services.perplexity import Perplexity from ml_model.services.pulid import Pulid from ml_model.services.photon import Photon +from ml_model.services.prunaai import Prunaai from ml_model.services.qwen import Qwen from ml_model.services.qwen_235B import Qwen_235B from ml_model.services.qwen_3_max_thinking import Qwen_3_Max_Thinking @@ -0,0 +1,63 @@ +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 replicate.exceptions import ModelError + +from messages.models import Message +from ml_model.exceptions import RequestBlocked, GenerationException +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 Elevenlabs(SimpleService): + TOKENS_COST = Decimal('2.490') + + @classmethod + def predict_price(cls, content: str, file_exists: bool, info: dict[str, Any]) -> Decimal | None: + duration = info['duration'] + return (cls.TOKENS_COST * duration).quantize(Decimal('0.1'), rounding='ROUND_UP') + + def calculate_price(self, duration: int) -> Decimal: + return (self.TOKENS_COST * 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), '.mp3'), + ) + if save: + return Message.objects.bulk_create([msg]) + return [msg] + + def make(self, input_message: Message, save: bool = True) -> list[Message]: + duration = input_message.info.pop('duration') + if (balance := PaymentPlanSelector(self.store.user).get_current_balance()) < ( + cost := self.calculate_price(duration) + ): + raise InsufficientBalance(balance, cost) + callback_data = { + 'prompt': input_message.content, + 'music_length_ms': duration * 1000, + **input_message.info, + } + start_time = time.time() + try: + audio = replicate_run('elevenlabs/music', callback_data) + except ModelError as exc: + if any(error in str(exc) for error in ('E005', 'E006', 'sexual')): + raise RequestBlocked + raise GenerationException from exc + process_time = timedelta(seconds=(time.time() - start_time)) + self.handle_invoice(input_message.content_object.model, duration=duration) + msgs = self.save_results(input_message.content, process_time, audio, save) + return msgs @@ -0,0 +1,69 @@ +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 replicate.exceptions import ModelError + +from messages.models import Message +from ml_model.exceptions import GenerationException, RequestBlocked +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 Prunaai(SimpleService): + TOKENS_COST = {'p-image': Decimal('2.5'), 'p-image-edit': Decimal('5'), 'flux-fast': Decimal('2.5')} + + def calculate_price(self, version: str) -> Decimal: + return self.TOKENS_COST[version] + + @classmethod + def predict_price(cls, content: str, file_exists: bool, info: dict[str, Any]) -> Decimal | None: + version = info['version'] + return cls.TOKENS_COST[version] + + def save_results(self, content: str, t: timedelta, image_url: str, save: bool = True) -> list[Message]: + msg = Message( + content=content, + content_object=self.store, + elapsed_time=t, + file=File(BytesIO(requests.get(image_url).content), '.png'), + ) + if save: + return Message.objects.bulk_create([msg]) + return [msg] + + def make(self, input_message: Message, save: bool = True) -> list[Message]: + version = input_message.info.get('version', 'p-image') + if input_message.file: + version = 'p-image-edit' + if (balance := PaymentPlanSelector(self.store.user).get_current_balance()) < self.TOKENS_COST[ + version + ]: + raise InsufficientBalance(balance, self.TOKENS_COST[version]) + callback_data = dict( + { + 'prompt': self.translate_prompt(input_message.content), + **input_message.info, + } + ) + if input_message.file: + callback_data.update({'images': [input_message.file.url]}) + start_time = time.time() + try: + images = replicate_run(f'prunaai/{version}', callback_data) + except ModelError as exc: + if any(error in str(exc) for error in ('E005', 'E006', 'sexual', 'NSFW')): + raise RequestBlocked + raise GenerationException from exc + process_time = timedelta(seconds=(time.time() - start_time)) + self.handle_invoice(input_message.content_object.model, version=version) + msgs = self.save_results(input_message.content, process_time, images, save) + return msgs