@@ -33,6 +33,7 @@ from ml_model.services.raifgpt import Raifgpt from ml_model.services.runway import Runway from ml_model.services.ray import Ray from ml_model.services.recraft import Recraft +from ml_model.services.reve import Reve from ml_model.services.sdxlemoji import Sdxlemoji from ml_model.services.stablediffusion import Stablediffusion from ml_model.services.upscaleai import Upscaleai @@ -0,0 +1,64 @@ +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.services.base import SimpleService +from ml_model.tasks import replicate_run + + +class Reve(SimpleService): + PRICE = { + 'create': Decimal('7.5'), + 'edit-fast': Decimal('3') + } + + def calculate_price(self, type: str) -> Decimal: + return self.PRICE[type].quantize(Decimal('0.1'), rounding='ROUND_UP') + + def save_results( + self, + prompt: str, + image: str, + time: timedelta, + save: bool = True, + ) -> list[Message]: + messages: list[Message] = [] + 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 = { + 'prompt': self.translate_prompt(input_message.content), + **input_message.info, + } + type = 'create' + if input_message.file: + kind = filetype.guess(input_message.file.read(20)) + mime = kind.mime if kind else 'application/octet-stream' + input_message.file.seek(0) + image = f'data:{mime};base64,{base64.b64encode(input_message.file.read()).decode("utf-8")}' + input_message.file.close() + callback_data.update({'image': image}) + type = 'edit-fast' + image = replicate_run(f'reve/{type}', callback_data) + process_time = timedelta(seconds=(time.time() - start_time)) + self.handle_invoice(input_message.content_object.model, type=type) + msgs = self.save_results(input_message.content, image, process_time, save) + return msgs