@@ -14,6 +14,7 @@ from ml_model.services.elevenlabs_music import Elevenlabs_Music from ml_model.services.epicphotogasm import Epicphotogasm from ml_model.services.flux import Flux from ml_model.services.flux_2 import Flux_2 +from ml_model.services.flux_3 import Flux_3 from ml_model.services.fluxkrea import Fluxkrea from ml_model.services.fluxlorafast import Fluxlorafast from ml_model.services.fluxproultra import Fluxproultra @@ -28,6 +29,7 @@ from ml_model.services.granite import Granite from ml_model.services.grok import Grok from ml_model.services.grok_4_1_fast import Grok_4_1_Fast from ml_model.services.grok_image import Grok_Image +from ml_model.services.grok_image_ultra import Grok_Image_Ultra from ml_model.services.grok_imagine_video import Grok_Imagine_Video from ml_model.services.hailuo import Hailuo from ml_model.services.hunyuan import Hunyuan @@ -0,0 +1,103 @@ +import base64 +import math +import time +from datetime import timedelta +from decimal import Decimal +from io import BytesIO +from typing import Any + +import filetype +import requests +from PIL import Image +from django.core.files import File +from django.core.files.images import get_image_dimensions + +from messages.models import Message +from ml_model.exceptions import FileExtensionNotSupported +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 Flux_3(SimpleService): + TOKENS_COST = { + 'input_mp': Decimal('1'), + 'output_mp': Decimal('7.5'), + } + + _MODEL = 'black-forest-labs/flux-2-klein-9b' + + def calculate_price(self, input_mp: int, output_mp: int) -> Decimal: + price = self.TOKENS_COST['input_mp'] * input_mp + self.TOKENS_COST['output_mp'] * output_mp + return price.quantize(Decimal('0.1'), rounding='ROUND_UP') + + @classmethod + def predict_price(cls, content: str, file_exists: bool, info: dict[str, Any]) -> Decimal | None: + if file_exists: + return None + return cls.TOKENS_COST['output_mp'].quantize(Decimal('0.1'), rounding='ROUND_UP') + + 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(messages) + return messages + + def make(self, input_message: Message, save: bool = True) -> list[Message]: + start_time = time.time() + output_megapixels = 1 + callback_data = { + 'prompt': input_message.content, + 'disable_safety_checker': False, + 'output_format': 'png', + 'output_megapixels': str(output_megapixels), + **input_message.info, + } + input_mp = 0 + if input_message.file: + file_bytes = input_message.file.read() + input_message.file.close() + kind = filetype.guess(file_bytes[:20]) + extension = kind.extension + if extension.upper() not in (extensions := ['JPG', 'JPEG', 'PNG', 'WEBP']): + raise FileExtensionNotSupported(extensions) + format = 'jpeg' if extension not in ('png', 'jpeg', 'webp') else extension + with Image.open(BytesIO(file_bytes)) as source_image: + normalized_image = source_image.convert('RGB') + with BytesIO() as buf: + normalized_image.save(buf, format=format) + file_width, file_height = get_image_dimensions(buf) + input_mp = math.ceil((file_width * file_height) / 1_000_000) + image = f'data:image/{format};base64,{base64.b64encode(buf.getvalue()).decode("utf-8")}' + normalized_image.close() + callback_data.update({'images': [image]}) + if (balance := PaymentPlanSelector(self.store.user).get_current_balance()) < ( + predicted := self.calculate_price(input_mp, output_megapixels) + ): + raise InsufficientBalance(balance, predicted) + images = replicate_run(self._MODEL, callback_data) + images = images if isinstance(images, list) else [images] + process_time = timedelta(seconds=(time.time() - start_time)) + self.handle_invoice( + input_message.content_object.model, + input_mp=input_mp, + output_mp=output_megapixels, + ) + msgs = self.save_results(input_message.content, images, process_time, save) + return msgs @@ -0,0 +1,94 @@ +import math +import time +from datetime import timedelta +from decimal import Decimal +from io import BytesIO +from typing import Any + +import filetype +import requests +from django.core.files import File +from django.core.files.images import get_image_dimensions + +from messages.models import Message +from ml_model.exceptions import FileExtensionNotSupported +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 Grok_Image_Ultra(SimpleService): + TOKENS_COST = { + '1k': Decimal('25'), + '2k': Decimal('35'), + 'input_mp': Decimal('5'), + } + + _MODEL = 'xai/grok-imagine-image-quality' + + def calculate_price(self, resolution: str, input_mp: int) -> Decimal: + price = self.TOKENS_COST[resolution] + self.TOKENS_COST['input_mp'] * input_mp + return price.quantize(Decimal('0.1'), rounding='ROUND_UP') + + @classmethod + def predict_price(cls, content: str, file_exists: bool, info: dict[str, Any]) -> Decimal | None: + if file_exists: + return None + resolution = info.get('resolution', '2k') + if resolution not in ('1k', '2k'): + return None + return cls.TOKENS_COST[resolution].quantize(Decimal('0.1'), rounding='ROUND_UP') + + def save_results( + self, + content: str, + image_url: str, + time: timedelta, + save: bool = True, + ) -> list[Message]: + message = Message( + content=content, + content_object=self.store, + elapsed_time=time, + file=File(BytesIO(requests.get(image_url).content), '.png'), + ) + if save: + return Message.objects.bulk_create([message]) + return [message] + + def make(self, input_message: Message, save: bool = True) -> list[Message]: + resolution = input_message.info.get('resolution', '2k') + callback_data = { + 'prompt': input_message.content, + **input_message.info, + } + input_mp = 0 + if input_message.file: + file_bytes = input_message.file.read() + kind = filetype.guess(file_bytes[:20]) + extension = kind.extension + if extension.upper() not in (extensions := ['JPG', 'JPEG', 'PNG', 'WEBP']): + raise FileExtensionNotSupported(extensions) + file_width, file_height = get_image_dimensions(BytesIO(file_bytes)) + if file_width and file_height: + input_mp = math.ceil((file_width * file_height) / 1_000_000) + else: + input_mp = 1 + callback_data.update({'image': input_message.file.url}) + input_message.file.close() + + if (balance := PaymentPlanSelector(self.store.user).get_current_balance()) < ( + predicted := self.calculate_price(resolution, input_mp) + ): + raise InsufficientBalance(balance, predicted) + + start_time = time.time() + image_url = replicate_run(self._MODEL, callback_data) + process_time = timedelta(seconds=(time.time() - start_time)) + self.handle_invoice( + input_message.content_object.model, + resolution=resolution, + input_mp=input_mp, + ) + return self.save_results(input_message.content, image_url, process_time, save) @@ -0,0 +1,59 @@ +# Generated by makemigration_payment_features on 2026-07-28 15:33 + +import math +from decimal import Decimal + +from django.db import migrations +from django.db.models import Max + + +def add_flux_3_payment_features(apps, schema_editor): + PaymentPlan = apps.get_model('payments', 'PaymentPlan') + PaymentPlanFeature = apps.get_model('payments', 'PaymentPlanFeature') + NeuronModel = apps.get_model('ml_model', 'NeuronModel') + + price = Decimal('9.5') + measurement_unit = 'file' + price_threshold = 0 + model = NeuronModel.objects.get(slug='flux_3') + category = model.category + + max_order_by_plan_id = { + row['plan_id']: row['max_order'] + for row in PaymentPlanFeature.objects.filter(model__category=category) + .values('plan_id') + .annotate(max_order=Max('order')) + } + + features = [] + for plan in PaymentPlan.objects.filter(price__gt=price_threshold): + quantity = math.floor(plan.tokens_per_plan / price) + max_order = max_order_by_plan_id.get(plan.pk) + next_order = (max_order if max_order is not None else -1) + 1 + max_order_by_plan_id[plan.pk] = next_order + features.append( + PaymentPlanFeature( + plan=plan, + model=model, + quantity=quantity, + measurement_unit=measurement_unit, + order=next_order, + ) + ) + PaymentPlanFeature.objects.bulk_create( + features, + update_conflicts=True, + update_fields=['quantity', 'measurement_unit'], + unique_fields=['plan', 'model'], + ) + + +class Migration(migrations.Migration): + + dependencies = [ + ('payments', '0031_remove_paymentmethod_attempts_and_more'), + ] + + operations = [ + migrations.RunPython(add_flux_3_payment_features, migrations.RunPython.noop), + ] @@ -0,0 +1,59 @@ +# Generated by makemigration_payment_features on 2026-07-28 15:34 + +import math +from decimal import Decimal + +from django.db import migrations +from django.db.models import Max + + +def add_grok_image_ultra_payment_features(apps, schema_editor): + PaymentPlan = apps.get_model('payments', 'PaymentPlan') + PaymentPlanFeature = apps.get_model('payments', 'PaymentPlanFeature') + NeuronModel = apps.get_model('ml_model', 'NeuronModel') + + price = Decimal('45') + measurement_unit = 'file' + price_threshold = 0 + model = NeuronModel.objects.get(slug='grok_image_ultra') + category = model.category + + max_order_by_plan_id = { + row['plan_id']: row['max_order'] + for row in PaymentPlanFeature.objects.filter(model__category=category) + .values('plan_id') + .annotate(max_order=Max('order')) + } + + features = [] + for plan in PaymentPlan.objects.filter(price__gt=price_threshold): + quantity = math.floor(plan.tokens_per_plan / price) + max_order = max_order_by_plan_id.get(plan.pk) + next_order = (max_order if max_order is not None else -1) + 1 + max_order_by_plan_id[plan.pk] = next_order + features.append( + PaymentPlanFeature( + plan=plan, + model=model, + quantity=quantity, + measurement_unit=measurement_unit, + order=next_order, + ) + ) + PaymentPlanFeature.objects.bulk_create( + features, + update_conflicts=True, + update_fields=['quantity', 'measurement_unit'], + unique_fields=['plan', 'model'], + ) + + +class Migration(migrations.Migration): + + dependencies = [ + ('payments', '0032_add_flux_3_payment_features'), + ] + + operations = [ + migrations.RunPython(add_grok_image_ultra_payment_features, migrations.RunPython.noop), + ]