@@ -112,6 +112,9 @@ REST_FRAMEWORK = { 'DEFAULT_FILTER_BACKENDS': ('django_filters.rest_framework.DjangoFilterBackend',), 'DEFAULT_SCHEMA_CLASS': 'drf_spectacular.openapi.AutoSchema', 'DEFAULT_PERMISSION_CLASSES': ('rest_framework.permissions.AllowAny',), + 'DEFAULT_RENDERER_CLASSES': ( + 'rest_framework.renderers.JSONRenderer', + ) } REST_USE_JWT = True SIMPLE_JWT = { @@ -2,13 +2,14 @@ import time from _decimal import Decimal from datetime import timedelta from io import BytesIO +from math import ceil 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 ml_model.tasks import replicate_run, fal_ai_run class Dalle(SimpleService): @@ -17,14 +18,14 @@ class Dalle(SimpleService): contains abstract method make, which makes a generation """ - PRICE = Decimal('2') + MP_PRICE = Decimal('1.5') - _CALLBACK = ( - 'bytedance/sdxl-lightning-4step:5599ed30703defd1d160a25a63321b4dec97101d98b4674bcc56e41f62f35637' - ) + _CALLBACK = 'fal-ai/flux/schnell' - def calculate_price(self, input_message: Message) -> Decimal: - price = input_message.info.get('num_outputs', 1) * self.PRICE + def calculate_price(self, input_message: Message, sizes: list[tuple]) -> Decimal: + price = Decimal('0') + for size in sizes: + price += Decimal(ceil(size[0] * size[1] / 1_000_000)) * self.MP_PRICE return price.quantize(Decimal('0.1'), rounding='ROUND_UP') def save_results(self, prompt: str, images: list, time: timedelta, save: bool = True) -> list[Message]: @@ -48,11 +49,17 @@ class Dalle(SimpleService): callback_data = dict( { 'prompt': translated_prompt, + 'image_size': { + 'width': input_message.info.pop('width', 1024), + 'height': input_message.info.pop('width', 1024) + }, **input_message.info, } ) - images = replicate_run(self._CALLBACK, callback_data) + result = fal_ai_run(self._CALLBACK, callback_data) + images = [img['url'] for img in result] + sizes = [(img['height'], img['width']) for img in result] process_time = timedelta(seconds=(time.time() - start_time)) - self.handle_invoice(input_message.content_object.model, input_message=input_message) + self.handle_invoice(input_message.content_object.model, input_message=input_message, sizes=sizes) msgs = self.save_results(input_message.content, images, process_time, save) return msgs @@ -2,6 +2,7 @@ import time from datetime import timedelta from decimal import Decimal from io import BytesIO +from math import ceil import requests from django.core.files import File @@ -11,7 +12,7 @@ from ml_model.models import ( NeuronModel, ) from ml_model.services.base import SimpleService -from ml_model.tasks import replicate_run +from ml_model.tasks import replicate_run, fal_ai_run class Flux(SimpleService): @@ -20,20 +21,15 @@ class Flux(SimpleService): contains abstract method make, which makes a generation """ - TOKENS_COST = { - 'flux-schnell': { - 'input_imgs': Decimal('3'), - }, - } + MP_PRICE = Decimal('1.5') # flux-schnell from fal.ai - def calculate_price(self, input_message: Message, version: str) -> Decimal: - price_map = self.TOKENS_COST[version] - price = price_map['input_imgs'] - if image_count := input_message.info.get('num_outputs'): - price = price * image_count + def calculate_price(self, input_message: Message, sizes: list[tuple]) -> Decimal: + price = Decimal('0') + for size in sizes: + price += Decimal(ceil(size[0] * size[1] / 1_000_000)) * self.MP_PRICE return price.quantize(Decimal('0.1'), rounding='ROUND_UP') - _CALLBACK_BASE = 'black-forest-labs/' + _CALLBACK = 'fal-ai/flux/schnell' @property def neuron_model(self): @@ -62,19 +58,22 @@ class Flux(SimpleService): def make(self, input_message: Message, save: bool = True) -> list[Message]: start_time = time.time() - version = input_message.info.get('version') + translated_prompt = self.translate_prompt(input_message.content) callback_data = dict( { - 'prompt': self.translate_prompt(input_message.content), + 'prompt': translated_prompt, + 'image_size': { + 'width': input_message.info.pop('width', 1024), + 'height': input_message.info.pop('width', 1024) + }, **input_message.info, } ) - runner = replicate_run( - f'{self._CALLBACK_BASE}{callback_data.get("version", "flux-schnell")}', - callback_data, - ) - images = runner if isinstance(runner, list) else [runner] + result = fal_ai_run(self._CALLBACK, callback_data) + images = [img['url'] for img in result] + sizes = [(img['height'], img['width']) for img in result] process_time = timedelta(seconds=(time.time() - start_time)) - self.handle_invoice(input_message.content_object.model, input_message=input_message, version=version) + self.handle_invoice(input_message.content_object.model, input_message=input_message, sizes=sizes) msgs = self.save_results(input_message.content, images, process_time, save) return msgs + @@ -3,6 +3,7 @@ import time from datetime import timedelta from decimal import Decimal from io import BytesIO +from math import ceil import filetype import requests @@ -14,7 +15,7 @@ from ml_model.models import ( NeuronModel, ) from ml_model.services.base import SimpleService -from ml_model.tasks import replicate_run +from ml_model.tasks import replicate_run, fal_ai_run class Fluxproultra(SimpleService): @@ -24,30 +25,27 @@ class Fluxproultra(SimpleService): """ TOKENS_COST = { - 'flux-dev': { - 'input_imgs': Decimal('7.5'), + 'flux/dev': { + 'mp': Decimal('12.5'), }, - 'flux-1.1-pro': { - 'input_imgs': Decimal('12.0'), + 'flux-pro/v1.1': { + 'mp': Decimal('20'), }, - 'flux-1.1-pro-ultra': { - 'input_imgs': Decimal('18'), + 'flux-pro/v1.1-ultra': { + 'img': Decimal('30'), # this model is billed by img quantity only }, } - def calculate_price(self, input_message: Message, version: str) -> Decimal: - price_map = self.TOKENS_COST[version] - price = price_map['input_imgs'] - if image_count := input_message.info.get('num_outputs'): - price = price * image_count + def calculate_price(self, input_message: Message, sizes: list[tuple] = None, version: str = '') -> Decimal: + if version == 'flux-pro/v1.1-ultra': + price = self.TOKENS_COST[version]['img'] * input_message.info.pop('num_images') + else: + price = Decimal('0') + for size in sizes: + price += Decimal(ceil(size[0] * size[1] / 1_000_000)) * self.TOKENS_COST[version]['mp'] return price.quantize(Decimal('0.1'), rounding='ROUND_UP') - inputs = [ - ModelInput(type=ModelInput.TypeChoices.TEXT, required=True), - ModelInput(type=ModelInput.TypeChoices.IMAGE), - ] - - _CALLBACK_BASE = 'black-forest-labs/' + _CALLBACK_BASE = 'fal-ai/' @property def neuron_model(self): @@ -80,22 +78,20 @@ class Fluxproultra(SimpleService): callback_data = dict( { 'prompt': self.translate_prompt(input_message.content), + 'image_size': { + 'width': input_message.info.pop('width', 1024), + 'height': input_message.info.pop('width', 1024) + }, **input_message.info, } ) - 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}) - runner = replicate_run( + result = fal_ai_run( f'{self._CALLBACK_BASE}{version}', callback_data, ) - images = runner if isinstance(runner, list) else [runner] + images = [img['url'] for img in result] + sizes = [(img['height'], img['width']) for img in result] process_time = timedelta(seconds=(time.time() - start_time)) - self.handle_invoice(input_message.content_object.model, input_message=input_message, version=version) + self.handle_invoice(input_message.content_object.model, input_message=input_message, sizes=sizes, version=version) msgs = self.save_results(input_message.content, images, process_time, save) return msgs @@ -8,7 +8,7 @@ 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 ml_model.tasks import replicate_run, fal_ai_run class Midjourney(SimpleService): @@ -17,14 +17,14 @@ class Midjourney(SimpleService): contains abstract method make, which makes a generation """ - _CALLBACK = 'minimax/image-01' - price = Decimal('2') + _CALLBACK = 'fal-ai/minimax/image-01' + price = Decimal('5') def __init__(self, store): super().__init__(store) def calculate_price(self, input_message: Message) -> Decimal: - price = input_message.info.get('number_of_images', 1) * self.price + price = input_message.info.get('num_images', 1) * self.price return price.quantize(Decimal('0.1'), rounding='ROUND_UP') def save_results( @@ -49,8 +49,9 @@ class Midjourney(SimpleService): translated_prompt = self.translate_prompt(input_message.content) activation_prompt = f'mdjrny-v4 style a highly detailed {translated_prompt}' callback_data = dict(prompt=activation_prompt, **input_message.info) - results = replicate_run(self._CALLBACK, callback_data) + result = fal_ai_run(self._CALLBACK, callback_data) + images = [img['url'] for img in result] process_time = timedelta(seconds=(time.time() - start_time)) self.handle_invoice(input_message.content_object.model, input_message=input_message) - msgs = self.save_results(input_message.content, results, process_time, save) + msgs = self.save_results(input_message.content, images, process_time, save) return msgs @@ -7,14 +7,8 @@ import requests from django.core.files import File from messages.models import Message -from ml_model.models import ( - ModelCategory, - ModelInput, - ModelParameter, - ModelVersion, -) from ml_model.services.base import SimpleService -from ml_model.tasks import replicate_run +from ml_model.tasks import replicate_run, fal_ai_run class Recraft(SimpleService): @@ -22,177 +16,147 @@ class Recraft(SimpleService): Recraft Service contains abstract method make, which makes a generation """ - - title = 'Recraft' - description = 'Нейросеть, способная генерировать картинки из вашего текста' - category = ModelCategory(title='Изображения', slug='images') - versions = [ - ModelVersion(name='Recraft V3', slug='recraft-v3'), - ModelVersion(name='Recraft V3 SVG', slug='recraft-v3-svg'), - ] - inputs = [ - ModelInput(type=ModelInput.TypeChoices.TEXT, required=True), - ] - parameters = [ - ModelParameter( - name='Ширина', - key='width', - type=ModelParameter.TypeChoices.INTRANGE, - values={'start': 1024, 'end': 2048, 'step': 128, 'default': 1024}, - ), - ModelParameter( - name='Высота', - key='height', - type=ModelParameter.TypeChoices.INTRANGE, - values={'start': 1024, 'end': 2048, 'step': 128, 'default': 1024}, - ), - ModelParameter( - name='Стиль', - key='style', - type=ModelParameter.TypeChoices.LIST, - values={ - 'availables': [ - 'любой', - 'реалистичное изображение', - 'цифровая иллюстрация', - 'пиксель-арт', - 'ручной рисунок', - 'зернистость', - 'детский рисунок', - '2D арт-постер', - 'ручной 3D', - 'контурный рисунок вручную', - 'гравировка в цвете', - '2D арт-постер 2', - 'черно-белое', - 'яркий свет', - 'HDR', - 'естественное освещение', - 'студийный портрет', - 'предпринимательство', - 'размытие движения', - ], - 'default': 'любой', - }, - ), - ModelParameter( - name='Стиль', - key='style', - type=ModelParameter.TypeChoices.LIST, - values={ - 'availables': [ - 'любой', - 'гравировка', - 'контурный рисунок', - 'схема', - 'линогравюра', - ], - 'default': 'любой', - }, - ), - ] - - payment_rules = { - versions[0].slug: Decimal('12'), - versions[1].slug: Decimal('24'), + PRICE = { + 'recraft-v3': Decimal('20'), + 'recraft-v3-svg': Decimal('40'), } - def _get_size(self, width: int, height: int) -> str: - available_sizes = ( - (1024, 1024), - (1365, 1024), - (1024, 1365), - (1536, 1024), - (1024, 1536), - (1820, 1024), - (1024, 1820), - (1024, 2048), - (2048, 1024), - (1434, 1024), - (1024, 1434), - (1024, 1280), - (1280, 1024), - (1024, 1707), - (1707, 1024), - ) - if width >= height: - size = min(available_sizes, key=lambda size: abs(width - size[0])) - else: - size = min(available_sizes, key=lambda size: abs(height - size[1])) - return f'{size[0]}x{size[1]}' + _CALLBACK = 'fal-ai/recraft/v3/text-to-image' def calculate_price(self, input_message: Message) -> Decimal: - return self.payment_rules[input_message.info.get('version', 'recraft-v3')] + return self.PRICE[input_message.info.get('version', 'recraft-v3')] def save_results( self, prompt: str, - image: str, + images: list, extension: 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), extension), + for image in images: + messages.append( + Message( + content_object=self.store, + elapsed_time=time, + content=prompt, + file=File(BytesIO(requests.get(image).content), extension), + ) ) - ) if save: return Message.objects.bulk_create(messages) return messages def make(self, input_message: Message, save: bool = True) -> list[Message]: styles = { - 'любой': 'any', - 'реалистичное изображение': 'realistic_image', - 'цифровая иллюстрация': 'digital_illustration', - 'пиксель-арт': 'digital_illustration/pixel_art', - 'ручной рисунок': 'digital_illustration/hand_drawn', - 'зернистость': 'digital_illustration/grain', - 'детский рисунок': 'digital_illustration/infantile_sketch', - '2D арт-постер': 'digital_illustration/2d_art_poster', - 'ручной 3D': 'digital_illustration/handmade_3d', - 'контурный рисунок вручную': 'digital_illustration/hand_drawn_outline', - 'гравировка в цвете': 'digital_illustration/engraving_color', - '2D арт-постер 2': 'digital_illustration/2d_art_poster_2', - 'черно-белое': 'realistic_image/b_and_w', - 'яркий свет': 'realistic_image/hard_flash', - 'HDR': 'realistic_image/hdr', - 'естественное освещение': 'realistic_image/natural_light', - 'студийный портрет': 'realistic_image/studio_portrait', - 'предпринимательство': 'realistic_image/enterprise', - 'размытие движения': 'realistic_image/motion_blur', - 'гравировка': 'engraving', - 'контурный рисунок': 'line_art', - 'схема': 'line_circuit', - 'линогравюра': 'linocut', + "любой": "any", + "реалистичное изображение": "realistic_image", + "цифровая иллюстрация": "digital_illustration", + "векторная иллюстрация": "vector_illustration", + "черно-белое": "realistic_image/b_and_w", + "яркий свет": "realistic_image/hard_flash", + "HDR": "realistic_image/hdr", + "естественное освещение": "realistic_image/natural_light", + "студийный портрет": "realistic_image/studio_portrait", + "предпринимательство": "realistic_image/enterprise", + "размытие движения": "realistic_image/motion_blur", + "вечернее освещение": "realistic_image/evening_light", + "выцветшая ностальгия": "realistic_image/faded_nostalgia", + "лесная жизнь": "realistic_image/forest_life", + "мистический натурализм": "realistic_image/mystic_naturalism", + "природные тона": "realistic_image/natural_tones", + "органическое спокойствие": "realistic_image/organic_calm", + "сияние реальной жизни": "realistic_image/real_life_glow", + "ретро-реализм": "realistic_image/retro_realism", + "ретро-снимок": "realistic_image/retro_snapshot", + "городская драма": "realistic_image/urban_drama", + "деревенский реализм": "realistic_image/village_realism", + "тёплое фольклорное": "realistic_image/warm_folk", + "пиксель-арт": "digital_illustration/pixel_art", + "ручной рисунок": "digital_illustration/hand_drawn", + "зернистость": "digital_illustration/grain", + "детский рисунок": "digital_illustration/infantile_sketch", + "2D арт-постер": "digital_illustration/2d_art_poster", + "ручной 3D": "digital_illustration/handmade_3d", + "контурный рисунок вручную": "digital_illustration/hand_drawn_outline", + "гравировка в цвете": "digital_illustration/engraving_color", + "2D арт-постер 2": "digital_illustration/2d_art_poster_2", + "антикварный стиль": "digital_illustration/antiquarian", + "яркое фэнтези": "digital_illustration/bold_fantasy", + "книга для детей": "digital_illustration/child_book", + "книги для детей": "digital_illustration/child_books", + "обложка": "digital_illustration/cover", + "штриховка": "digital_illustration/crosshatch", + "цифровая гравюра": "digital_illustration/digital_engraving", + "экспрессионизм": "digital_illustration/expressionism", + "детали от руки": "digital_illustration/freehand_details", + "зернистость 20%": "digital_illustration/grain_20", + "графическая насыщенность": "digital_illustration/graphic_intensity", + "комиксы с резкими линиями": "digital_illustration/hard_comics", + "длинная тень": "digital_illustration/long_shadow", + "современный фольклор": "digital_illustration/modern_folk", + "многоцветная": "digital_illustration/multicolor", + "неоновое спокойствие": "digital_illustration/neon_calm", + "нуар": "digital_illustration/noir", + "ностальгический пастель": "digital_illustration/nostalgic_pastel", + "контурные детали": "digital_illustration/outline_details", + "пастельный градиент": "digital_illustration/pastel_gradient", + "пастельный набросок": "digital_illustration/pastel_sketch", + "поп-арт": "digital_illustration/pop_art", + "поп-ренессанс": "digital_illustration/pop_renaissance", + "стрит-арт": "digital_illustration/street_art", + "набросок с планшета": "digital_illustration/tablet_sketch", + "городское сияние": "digital_illustration/urban_glow", + "городской скетчинг": "digital_illustration/urban_sketching", + "ванильные мечты": "digital_illustration/vanilla_dreams", + "книга для молодежи": "digital_illustration/young_adult_book", + "книга для молодежи 2": "digital_illustration/young_adult_book_2", + "яркий контур": "vector_illustration/bold_stroke", + "химическая графика": "vector_illustration/chemistry", + "цветной трафарет": "vector_illustration/colored_stencil", + "контурный поп-арт": "vector_illustration/contour_pop_art", + "космическая графика": "vector_illustration/cosmics", + "вырез": "vector_illustration/cutout", + "депрессивная стилистика": "vector_illustration/depressive", + "редакционная графика": "vector_illustration/editorial", + "эмоциональный flat": "vector_illustration/emotional_flat", + "инфографика": "vector_illustration/infographical", + "контур маркером": "vector_illustration/marker_outline", + "мозаика": "vector_illustration/mosaic", + "наивная векторная": "vector_illustration/naivector", + "округлый flat": "vector_illustration/roundish_flat", + "сегментированные цвета": "vector_illustration/segmented_colors", + "резкий контраст": "vector_illustration/sharp_contrast", + "тонкий контур": "vector_illustration/thin", + "векторная фотография": "vector_illustration/vector_photo", + "яркие формы": "vector_illustration/vivid_shapes", + "гравировка": "vector_illustration/engraving", + "линейный рисунок": "vector_illustration/line_art", + "схема": "vector_illustration/line_circuit", + "линогравюра": "vector_illustration/linocut", } start_time = time.time() extension = ( - '.svg' if input_message.info.get('version', 'recraft-v3') == self.versions[1].slug else '.png' - ) - size = self._get_size( - input_message.info.pop('width', 1024), - input_message.info.pop('height', 1024), + '.svg' if input_message.info.get('version', 'recraft-v3') == 'recraft-v3-svg' else '.png' ) - callback_url = f'recraft-ai/{input_message.info.get("version", "recraft-v3")}' callback_data = dict( { 'prompt': ( 'The subject should be fully clothed, in a neutral or formal style. ' f'Input description: {self.translate_prompt(input_message.content)}' ), - 'size': size, + 'image_size': { + 'width': input_message.info.pop('width', 1024), + 'height': input_message.info.pop('width', 1024) + }, 'style': styles.get(input_message.info.pop('style'), 'любой'), **input_message.info, } ) - image = replicate_run(callback_url, callback_data) + result = fal_ai_run(self._CALLBACK, callback_data) + images = [img['url'] for img in result] process_time = timedelta(seconds=(time.time() - start_time)) self.handle_invoice(input_message.content_object.model, input_message) - message = self.save_results(input_message.content, image, extension, process_time, save) + message = self.save_results(input_message.content, images, extension, process_time, save) return message @@ -1,17 +1,16 @@ import logging import time -import uuid from datetime import timedelta from decimal import Decimal from io import BytesIO +from math import ceil -import httpx -from django.conf import settings +import requests from django.core.files import File from messages.models import Message from ml_model.services.base import SimpleService -from poller.models import Proxy +from ml_model.tasks import fal_ai_run logger = logging.getLogger(__name__) @@ -22,43 +21,37 @@ class Stablediffusion(SimpleService): contains abstract method make, which makes a generation """ - MODELS = ['sd3', 'sd3-turbo', 'sd3-medium'] + MODELS = ['sd3', 'sd3-medium'] MODELS_LINKS = { - 'sd3': 'stable-diffusion-3.5-large', - 'sd3-turbo': 'stable-diffusion-3.5-large-turbo', - 'sd3-medium': 'stable-diffusion-3.5-medium' + 'sd3': 'stable-diffusion-v35-large', + 'sd3-medium': 'stable-diffusion-v35-medium' } + PRICES = { + 'stable-diffusion-v35-large': Decimal('32.5'), + 'stable-diffusion-v35-medium': Decimal('10'), + } + _CALLBACK = 'fal-ai/' - def calculate_price(self, input_message: Message) -> Decimal: - if input_message.info.get('version') == 'sd3': - return Decimal('32.5') - elif input_message.info.get('version') == 'sd3-turbo': - return Decimal('20') - elif input_message.info.get('version') == 'sd3-medium': - return Decimal('17.5') + def calculate_price(self, input_message: Message, sizes: list[tuple], version: str) -> Decimal: + price = Decimal('0') + for size in sizes: + price += Decimal(ceil(size[0] * size[1] / 1_000_000)) * self.PRICES[version] + return price.quantize(Decimal('0.1'), rounding='ROUND_UP') - def save_results( - self, - input_prompt: str, - link: str, - t: timedelta, - save: bool = True, - ) -> list[Message]: - out: list[Message] = [] - out.append( - Message( - content_object=self.store, - elapsed_time=t, - content=input_prompt, - file=File( - BytesIO(httpx.get(link).content), - f'{uuid.uuid4()}.png', - ), + 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(out) - return out + return Message.objects.bulk_create(messages) + return messages def make(self, input_message: Message, save: bool = True) -> list[Message]: start_time = time.time() @@ -67,33 +60,14 @@ class Stablediffusion(SimpleService): translated_prompt = self.translate_prompt(input_message.content) callback_data = { 'prompt': translated_prompt, - 'aspect_ratio': input_message.info.get('aspect_ratio', '1:1'), - 'output_quality': 100, + 'image_size': 'square_hd', 'output_format': 'png', + **input_message.info, } - link = '' - for proxy in Proxy.objects.all(): - with httpx.Client( - headers={ - 'Authorization': f'Bearer {settings.REPLICATE_API_KEY}', - 'Prefer': 'wait', - 'Content-Type': 'application/json', - }, - timeout=600, - proxy=f'{proxy.protocol}://{proxy.address}', - ) as client: - result = client.post( - f'https://api.replicate.com/v1/models/stability-ai/{model_name}/predictions', - json={'input': callback_data}, - ) - while result.json()['status'] not in ('succeeded', 'failed', 'canceled'): - result = client.get(result.json()['urls']['get']) - if result.json()['status'] in ('failed', 'canceled'): - logger.error(result.json()['logs']) - raise Exception('No answer from Stable Diffusion, please retry later') - link = result.json()['output'][0] - - process_time = timedelta(seconds=(time.time() - start_time)) - self.handle_invoice(input_message.content_object.model, input_message) - msgs = self.save_results(input_message.content, link, process_time, save) - return msgs + result = fal_ai_run(f'{self._CALLBACK}{model_name}', callback_data) + images = [img['url'] for img in result] + sizes = [(img['height'], img['width']) for img in result] + process_time = timedelta(seconds=(time.time() - start_time)) + self.handle_invoice(input_message.content_object.model, input_message, sizes, model_name) + msgs = self.save_results(input_message.content, images, process_time, save) + return msgs @@ -2,6 +2,7 @@ import base64 import json import logging import re +import time # import uuid from io import BytesIO @@ -17,7 +18,7 @@ from deepl.translator import TextResult from requests import Response from backend import settings -from ml_model.exceptions import DeploymentDisabled +from ml_model.exceptions import DeploymentDisabled, ModelTimeoutError from ml_model.utils import count_openrouter_tokens from poller.models import Proxy @@ -166,6 +167,29 @@ def openrouter_run(version: str, messages: list, callback_data: dict, model_name raise Exception(f'No answer from {model_name}, please retry later') +@shared_task +def fal_ai_run(model, payload): + requests_number = 0 + client = httpx.Client( + base_url="https://queue.fal.run", + headers={"Authorization": f"Key {settings.FAL_API_KEY}"}, + timeout=600, + ) + result = client.post( + model, + json=payload + ).json() + while True: + status = client.get(result['status_url']).json() + if status.get('status') == 'COMPLETED': + break + requests_number += 1 + if requests_number == 271: + raise ModelTimeoutError + time.sleep(1 / 3) + return client.get(result['response_url']).json()['images'] + + @shared_task def upscale_run(payload: dict[str, tuple[str, IO]]) -> list[str]: content = requests.post(