@@ -1134,6 +1134,14 @@ msgstr "1080р разрешение не поддерживается для See msgid "3K output is not supported for this model" msgstr "3К разрешение не поддерживается для этой модели" +#: ml_model/services/flux_3.py:87 +msgid "Draft mode is only available at 720p" +msgstr "Режим draft доступен только при 720p" + +#: ml_model/services/FileService.py +msgid "The attached video must be at most %(max_seconds)d seconds" +msgstr "Прикреплённое видео должно быть не длиннее %(max_seconds)d секунд" + #: ml_model/services/upscaleai.py:124 msgid "No image given for improving" msgstr "Нет изображения для улучшения" @@ -1,3 +1,4 @@ +import json import re import subprocess import zipfile @@ -123,48 +124,48 @@ class FileProcessingService: return voice.file -class ImageFileProcessingService: - ALLOWED_EXTENSIONS = ['PNG', 'JPG', 'JPEG', 'WEBP'] +class MediaFileProcessor: + EXTENSIONS = [] - def __init__(self, image: FieldFile) -> None: - self.image = image + def __init__(self, media_file: FieldFile) -> None: + self._media_file = media_file @staticmethod - def __reset_image(func): + def _reset_media_file(func): @wraps(func) def wrapper(self, *args, **kwargs): try: return func(self, *args, **kwargs) finally: - self.image.seek(0) + if not self._media_file.closed: + self._media_file.seek(0) return wrapper - @__reset_image + @_reset_media_file def get_bytes(self, size: int | None = None) -> bytes: - return self.image.read(size) + return self._media_file.read(size) - def get_kind(self, file_bytes: bytes): + @staticmethod + def get_kind(file_bytes): kind = filetype.guess(file_bytes) - if not kind: - raise CorruptedFileError - if kind.extension.upper() not in self.ALLOWED_EXTENSIONS: - raise FileExtensionNotSupported(self.ALLOWED_EXTENSIONS) return kind - @__reset_image - def get_dimensions(self, max_pixels: int) -> tuple[int, int]: + +class ImageFileProcessor(MediaFileProcessor): + EXTENSIONS = ['PNG', 'JPG', 'JPEG', 'WEBP'] + + @MediaFileProcessor._reset_media_file + def get_dimensions(self) -> tuple[int | None, int | None]: try: - w, h = get_image_dimensions(self.image) - if not (w and h): - raise CorruptedFileError - if w * h > max_pixels: - raise ImageTooLargeError(max_pixels) - return w, h - except Image.DecompressionBombError: - raise ImageTooLargeError(max_pixels) + return get_image_dimensions(self._media_file) + except Image.DecompressionBombError as exc: + if (max_pixels := Image.MAX_IMAGE_PIXELS) is None: + raise + raise ImageTooLargeError(max_pixels) from exc - def get_normalized_image(self, file_bytes: bytes) -> BytesIO: + @staticmethod + def get_normalized_image(file_bytes: bytes) -> BytesIO: normalized_image = BytesIO(file_bytes) with Image.open(normalized_image) as source_image: img = source_image.convert('RGBA') @@ -173,3 +174,59 @@ class ImageFileProcessingService: img.close() normalized_image.seek(0) return normalized_image + + +class VideoFileProcessor(MediaFileProcessor): + EXTENSIONS = ['MP4'] + + def get_duration(self) -> float: + result = subprocess.run( + [ + 'ffprobe', + '-v', + 'error', + '-show_entries', + 'format=duration', + '-of', + 'json', + '-', + ], + input=self.get_bytes(), + capture_output=True, + check=False, + ) + if result.returncode != 0: + raise CorruptedFileError + duration = float(json.loads(result.stdout).get('format', {}).get('duration') or 0) + if duration <= 0: + raise CorruptedFileError + return duration + + +class MediaFileValidator: + @staticmethod + def validate_kind(kind, allowed_extensions: list[str]) -> None: + if not kind: + raise CorruptedFileError + if kind.extension.upper() not in allowed_extensions: + raise FileExtensionNotSupported(allowed_extensions) + + +class ImageFileValidator(MediaFileValidator): + @staticmethod + def validate_dimensions(w: int | None, h: int | None, *, max_pixels: int | None) -> None: + if not (w and h): + raise CorruptedFileError + if max_pixels: + if w * h > max_pixels: + raise ImageTooLargeError(max_pixels) + + +class VideoFileValidator(MediaFileValidator): + @staticmethod + def validate_duration(duration: float, *, max_seconds: int) -> None: + if duration > max_seconds: + raise InvalidParameterError( + _('The attached video must be at most %(max_seconds)d seconds') + % {'max_seconds': max_seconds} + ) @@ -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 @@ -0,0 +1,119 @@ +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 django.utils.translation import gettext as _ + +from messages.models import Message +from ml_model.exceptions import InvalidParameterError +from ml_model.services.FileService import ( + ImageFileProcessor, + MediaFileProcessor, + MediaFileValidator, + VideoFileProcessor, + VideoFileValidator, +) +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): + MAX_START_VIDEO_SECONDS = 15 + TOKENS_COST = { + 't2v_i2v': { + '720p': Decimal('51'), # $0.17 / sec + '1080p': Decimal('87'), # $0.29 / sec + 'draft': Decimal('18'), # $0.06 / sec + }, + 'v2v': { + '720p': Decimal('123'), # $0.41 / sec + '1080p': Decimal('159'), # $0.53 / sec + 'draft': Decimal('36'), # $0.12 / sec + }, + } + + @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', '720p') + duration = int(info.get('duration', 5)) + draft = bool(info.get('draft', False)) + if resolution not in ('720p', '1080p'): + return None + if draft and resolution != '720p': + return None + key = 'draft' if draft else resolution + return (cls.TOKENS_COST['t2v_i2v'][key] * duration).quantize(Decimal('0.1'), rounding='ROUND_UP') + + def calculate_price(self, variant: str, resolution: str, duration: int, draft: bool) -> Decimal: + key = 'draft' if draft else resolution + return (self.TOKENS_COST[variant][key] * 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), '.mp4'), + ) + if save: + return Message.objects.bulk_create([msg]) + return [msg] + + def make(self, input_message: Message, save: bool = True) -> list[Message]: + resolution = input_message.info.pop('resolution', '720p') + duration = int(input_message.info.pop('duration', 5)) + draft = bool(input_message.info.get('draft', False)) + input_message.info.pop('safety_tolerance', None) + if draft and resolution != '720p': + raise InvalidParameterError(_('Draft mode is only available at 720p')) + variant = 't2v_i2v' + callback_data = dict( + { + 'prompt': input_message.content, + 'resolution': resolution, + 'duration': str(duration), + **input_message.info, + 'safety_tolerance': 2, + } + ) + if file := input_message.file: + processor = MediaFileProcessor(file) + kind = processor.get_kind(processor.get_bytes(50)) + allowed_extensions = ImageFileProcessor.EXTENSIONS + VideoFileProcessor.EXTENSIONS + MediaFileValidator.validate_kind(kind, allowed_extensions) + if kind.extension.upper() in VideoFileProcessor.EXTENSIONS: + video_processor = VideoFileProcessor(file) + VideoFileValidator.validate_duration( + video_processor.get_duration(), + max_seconds=self.MAX_START_VIDEO_SECONDS, + ) + callback_data.update({'start_video': file.url}) + variant = 'v2v' + else: + callback_data.update({'images': [file.url]}) + file.close() + if (balance := PaymentPlanSelector(self.store.user).get_current_balance()) < ( + predicted := self.calculate_price(variant, resolution, duration, draft) + ): + raise InsufficientBalance(balance, predicted) + start_time = time.time() + video = replicate_run('black-forest-labs/flux-3', callback_data) + process_time = timedelta(seconds=(time.time() - start_time)) + self.handle_invoice( + input_message.content_object.model, + variant=variant, + resolution=resolution, + duration=duration, + draft=draft, + ) + msgs = self.save_results(input_message.content, process_time, video, save) + return msgs @@ -12,7 +12,7 @@ from messages.models import Message from ml_model.adapters.bytedance_model_ark import BytedanceContentType from ml_model.exceptions import InvalidParameterError from ml_model.exceptions import ModelVersionNotAvailable -from ml_model.services.FileService import ImageFileProcessingService +from ml_model.services.FileService import ImageFileProcessor, ImageFileValidator from ml_model.services.base import SimpleService from ml_model.tasks import bytedance_model_ark_run from payments.exceptions.insufficient_balance import InsufficientBalance @@ -94,11 +94,12 @@ class Seedream(SimpleService): **input_message.info, } if image := input_message.file: - image_processor = ImageFileProcessingService(image) - file_bytes = image_processor.get_bytes(20) - image_processor.get_kind(file_bytes) - image_processor.get_dimensions(max_pixels=36000000) - image_processor.image.close() + image_processor = ImageFileProcessor(image) + kind = image_processor.get_kind(image_processor.get_bytes(20)) + ImageFileValidator.validate_kind(kind, image_processor.EXTENSIONS) + width, height = image_processor.get_dimensions() + ImageFileValidator.validate_dimensions(width, height, max_pixels=36000000) + image.close() callback_data.update({'image': image.url}) images = bytedance_model_ark_run( @@ -0,0 +1,66 @@ +# Generated by makemigration_payment_features on 2026-08-13 18:42 + +import math +import logging +from decimal import Decimal + +from django.db import migrations +from django.db.models import Max + +logger = logging.getLogger(__name__) + + +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('90') + measurement_unit = 'file' + price_threshold = 0 + try: + model = NeuronModel.objects.get(slug='flux_3') + except NeuronModel.DoesNotExist: + logger.warning('The model %s was not found. Migration will do nothing.', 'flux_3') + return + 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', '0035_paymentplanuserinfo_recovery_fields'), + ] + + operations = [ + migrations.RunPython(add_flux_3_payment_features, migrations.RunPython.noop), + ]