@@ -14,7 +14,6 @@ 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 @@ -1,54 +1,42 @@ +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.models import ( - NeuronModel, -) +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(SimpleService): - """ - Flux Service - contains abstract method make, which makes a generation - """ - TOKENS_COST = { - 'flux-schnell': { - 'input_imgs': Decimal('3'), - }, + 'input_mp': Decimal('1'), + 'output_mp': Decimal('7.5'), } - OPTIMIZATION_PROMPT = """ - Clean composition with clear subject hierarchy. Soft natural lighting, accurate proportions, and coherent geometry. - Sharp key details with minimal visual noise and artifacts. - """ + _MODEL = 'black-forest-labs/flux-2-klein-9b' - 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_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: - price = cls.TOKENS_COST['flux-schnell']['input_imgs'] * info.get('num_outputs', 1) - return price.quantize(Decimal('0.1'), rounding='ROUND_UP') - - _CALLBACK_BASE = 'black-forest-labs/' - - @property - def neuron_model(self): - return NeuronModel.objects.get(title='Flux') + if file_exists: + return None + return cls.TOKENS_COST['output_mp'].quantize(Decimal('0.1'), rounding='ROUND_UP') def save_results( self, @@ -73,22 +61,43 @@ class Flux(SimpleService): def make(self, input_message: Message, save: bool = True) -> list[Message]: start_time = time.time() - version = 'flux-schnell' - user_prompt = self.translate_prompt(input_message.content) - callback_data = dict( - { - 'prompt': f"{user_prompt}\n{self.OPTIMIZATION_PROMPT}", - 'go_fast': False, - 'output_quality': 100, - **input_message.info, - } - ) - runner = replicate_run( - f'{self._CALLBACK_BASE}{version}', - callback_data, - ) - images = runner if isinstance(runner, list) else [runner] + 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_message=input_message, version=version) + 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 @@ -1,103 +0,0 @@ -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,63 @@ +# Generated by makemigration_payment_features on 2026-07-28 17:27 + +import math +from decimal import Decimal + +from django.db import migrations +from django.db.models import Max + + +def add_flux_payment_features(apps, schema_editor): + PaymentPlan = apps.get_model('payments', 'PaymentPlan') + PaymentPlanFeature = apps.get_model('payments', 'PaymentPlanFeature') + NeuronModel = apps.get_model('ml_model', 'NeuronModel') + + try: + model = NeuronModel.objects.get(slug='flux') + except NeuronModel.DoesNotExist: + return + + price = Decimal('9.5') + measurement_unit = 'file' + price_threshold = 0 + 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', '0033_add_grok_image_ultra_payment_features'), + ] + + operations = [ + migrations.RunPython(add_flux_payment_features, migrations.RunPython.noop), + ]