@@ -334,6 +334,7 @@ GOOGLE_API_KEY = env.str('GOOGLE_API_KEY', default='defaultapikey') SERPER_API_KEY = env.str('SERPER_API_KEY', 'defaultapikey') FLUX_API_KEY = env.str('FLUX_API_KEY', 'defaultapikey') OPENROUTER_API_KEY = env.str('OPENROUTER_API_KEY', 'defaultapikey') +BYTEDANCE_MODEL_ARK_API_KEY = env.str('BYTEDANCE_MODEL_ARK_API_KEY', 'defaultapikey') FAL_API_KEY = env.str('FAL_API_KEY', 'defaultapikey') YANDEX_CLOUD_API_KEY = env.str('YANDEX_CLOUD_API_KEY', 'defaultapikey') @@ -0,0 +1 @@ +from ml_model.adapters.bytedance_model_ark import BytedanceModelArkAdapter \ No newline at end of file @@ -0,0 +1,317 @@ +from enum import StrEnum +import logging +import time +from typing import Any, TypeAlias, TypedDict + +import httpx + +from backend import settings +from ml_model.exceptions import ( + FileExtensionNotSupported, + GenerationException, + RequestBlocked, +) +from poller.models import Proxy + +logger = logging.getLogger(__name__) + + +class BytedanceContentType(StrEnum): + CHAT = 'chat' + IMAGE = 'image' + VIDEO = 'video' + + +class BytedanceFinishReason(StrEnum): + STOP = 'stop' + CONTENT_FILTER = 'content_filter' + LENGTH = 'length' + + +class BytedanceVideoTaskStatus(StrEnum): + QUEUED = 'queued' + RUNNING = 'running' + CANCELLED = 'cancelled' + SUCCEEDED = 'succeeded' + FAILED = 'failed' + EXPIRED = 'expired' + + +class BytedanceUsage(TypedDict, total=False): + prompt_tokens: int + completion_tokens: int + + +class BytedanceChatMessage(TypedDict, total=False): + content: str + reasoning_content: str + + +class BytedanceChatChoice(TypedDict, total=False): + finish_reason: str + message: BytedanceChatMessage + + +class BytedanceChatResponse(TypedDict, total=False): + choices: list[BytedanceChatChoice] + usage: BytedanceUsage + error: dict[str, Any] + + +class BytedanceVideoTaskResponse(TypedDict, total=False): + id: str + task_id: str + status: str + video_url: str + error: dict[str, Any] + + +RunChatResult: TypeAlias = tuple[str, int, int] +RunImageResult: TypeAlias = list[str] +RunVideoResult: TypeAlias = tuple[str, int] +BytedanceRunResult: TypeAlias = RunChatResult | RunImageResult | RunVideoResult + + +class BytedanceModelArkAdapter: + BASE_URL = 'https://ark.ap-southeast.bytepluses.com/api/v3/' + CONTENT_TYPE_TO_ENDPOINT = { + BytedanceContentType.CHAT: 'chat/completions', + BytedanceContentType.IMAGE: 'images/generations', + BytedanceContentType.VIDEO: 'contents/generations/tasks', + } + IMAGE_EXTENSIONS = {'jpg', 'jpeg', 'png', 'webp'} + VIDEO_EXTENSIONS = {'mp4'} + AUDIO_EXTENSIONS = {'mp3', 'wav', 'flac', 'aac', 'ogg', 'm4a'} + USER_ATTACHMENT_EXTENSIONS = ('JPG', 'JPEG', 'PNG', 'WEBP', 'MP4') + VIDEO_TASK_ENDPOINT = 'contents/generations/tasks' + VIDEO_POLL_ATTEMPTS_LIMIT = 271 + VIDEO_POLL_DELAY_SECONDS = 1 / 3 + + @classmethod + def validate_user_attachment_extension(cls, file_extension: str | None) -> None: + if not file_extension: + raise FileExtensionNotSupported(cls.USER_ATTACHMENT_EXTENSIONS) + + extension = file_extension.lower().strip('.') + if extension in cls.AUDIO_EXTENSIONS: + # NOTE: Пока не добавляем аудио, но технически возможно + raise FileExtensionNotSupported(cls.USER_ATTACHMENT_EXTENSIONS) + if extension not in cls.IMAGE_EXTENSIONS | cls.VIDEO_EXTENSIONS: + raise FileExtensionNotSupported(cls.USER_ATTACHMENT_EXTENSIONS) + + @classmethod + def get_user_attachment_type(cls, file_extension: str | None) -> str: + cls.validate_user_attachment_extension(file_extension) + extension = (file_extension or '').lower().strip('.') + if extension in cls.VIDEO_EXTENSIONS: + return 'video_url' + return 'image_url' + + @classmethod + def _raise_by_error_payload(cls, data: dict[str, Any], choices: list[dict[str, Any]]) -> None: + choice_reasons = {str(choice.get('finish_reason', '')).lower() for choice in choices} + + if BytedanceFinishReason.CONTENT_FILTER in choice_reasons: + raise RequestBlocked + + raise GenerationException + + + @classmethod + def _extract_chat_answer( + cls, + data: BytedanceChatResponse, + include_reasoning: bool = False, + ) -> str: + choices = data.get('choices') or [] + stop_choices = [ + choice + for choice in choices + if choice.get('finish_reason') == BytedanceFinishReason.STOP + and choice.get('message') + and choice['message'].get('content') is not None + ] + if stop_choices: + content = ','.join(str(choice['message']['content']) for choice in stop_choices) + # TODO: включить после разделения reasoning и content в хранении + if include_reasoning: + reasoning = ','.join( + str(reasoning) + for choice in stop_choices + if choice.get('message') + and (reasoning := choice['message'].get('reasoning_content')) is not None + ) + if reasoning and content: + return f'**Рассуждение:**\n\n{reasoning}\n\n**Основная мысль:**\n\n{content}' + if reasoning: + return reasoning + return content + + cls._raise_by_error_payload(data, choices) + + @classmethod + def run( + cls, + model: str, + callback_data: dict[str, Any], + content_type: BytedanceContentType = BytedanceContentType.CHAT, + messages: list[dict[str, Any]] | None = None, + include_reasoning: bool = False, + ) -> BytedanceRunResult: + if content_type == BytedanceContentType.IMAGE: + return cls._generate_image(model=model, callback_data=callback_data) + if content_type == BytedanceContentType.VIDEO: + return cls._generate_video(model=model, callback_data=callback_data) + return cls._run_chat( + model=model, + callback_data=callback_data, + messages=messages, + include_reasoning=include_reasoning, + ) + + @classmethod + def _run_chat( + cls, + model: str, + callback_data: dict[str, Any], + messages: list[dict[str, Any]] | None = None, + include_reasoning: bool = False, + ) -> RunChatResult: + payload = dict(callback_data) + payload['model'] = model + payload['messages'] = messages or [] + for proxy in Proxy.objects.all(): + with httpx.Client( + base_url=cls.BASE_URL, + headers={ + 'Authorization': f'Bearer {settings.BYTEDANCE_MODEL_ARK_API_KEY}', + 'Content-Type': 'application/json', + }, + proxy=f'{proxy.protocol}://{proxy.address}', + timeout=600, + ) as client: + resp = client.post(cls.CONTENT_TYPE_TO_ENDPOINT[BytedanceContentType.CHAT], json=payload) + if resp.status_code >= 400: + logger.error( + 'Bytedance request failed status=%s route=%s body=%s', + resp.status_code, + cls.CONTENT_TYPE_TO_ENDPOINT[BytedanceContentType.CHAT], + resp.text, + ) + data: BytedanceChatResponse = resp.json() + try: + answer = cls._extract_chat_answer( + data=data, + include_reasoning=include_reasoning, + ) + usage = data.get('usage') or {} + input_tokens = int(usage.get('prompt_tokens', 0)) + output_tokens = int(usage.get('completion_tokens', 0)) + + return answer, input_tokens, output_tokens + except (RequestBlocked, GenerationException): + raise + except Exception as exc: + logger.error('Invalid Bytedance Ark response: %s', data) + raise GenerationException from exc + + raise GenerationException + + @classmethod + def _generate_video( + cls, + model: str, + callback_data: dict[str, Any], + ) -> RunVideoResult: + if task_id := callback_data.get('task_id'): + return cls._retrieve_video_task(task_id=str(task_id)) + task_data = cls._create_video_task(model=model, callback_data=callback_data) + created_task_id = task_data.get('id') or task_data.get('task_id') + if not created_task_id: + raise GenerationException + return cls._wait_video_task_result(task_id=str(created_task_id)) + + @classmethod + def _generate_image(cls, model: str, callback_data: dict[str, Any]) -> list[str]: + payload = dict(callback_data) + payload['model'] = model + for proxy in Proxy.objects.all(): + with httpx.Client( + base_url=cls.BASE_URL, + headers={ + 'Authorization': f'Bearer {settings.BYTEDANCE_MODEL_ARK_API_KEY}', + 'Content-Type': 'application/json', + }, + proxy=f'{proxy.protocol}://{proxy.address}', + timeout=600, + ) as client: + resp = client.post(cls.CONTENT_TYPE_TO_ENDPOINT[BytedanceContentType.IMAGE], json=payload) + data = resp.json() + if image_data := data.get('data'): + urls = [item.get('url') for item in image_data if isinstance(item, dict) and item.get('url')] + if urls: + return urls + + raise GenerationException + + @classmethod + def _create_video_task(cls, model: str, callback_data: dict[str, Any]) -> BytedanceVideoTaskResponse: + payload = dict(callback_data) + payload['model'] = model + for proxy in Proxy.objects.all(): + with httpx.Client( + base_url=cls.BASE_URL, + headers={ + 'Authorization': f'Bearer {settings.BYTEDANCE_MODEL_ARK_API_KEY}', + 'Content-Type': 'application/json', + }, + proxy=f'{proxy.protocol}://{proxy.address}', + timeout=600, + ) as client: + route = cls.CONTENT_TYPE_TO_ENDPOINT[BytedanceContentType.VIDEO] + resp = client.post(route, json=payload) + try: + data: BytedanceVideoTaskResponse = resp.json() + except Exception: + raise GenerationException + if data.get('id') or data.get('task_id'): + return data + + raise GenerationException + + @classmethod + def _retrieve_video_task(cls, task_id: str) -> BytedanceVideoTaskResponse: + for proxy in Proxy.objects.all(): + with httpx.Client( + base_url=cls.BASE_URL, + headers={ + 'Authorization': f'Bearer {settings.BYTEDANCE_MODEL_ARK_API_KEY}', + 'Content-Type': 'application/json', + }, + proxy=f'{proxy.protocol}://{proxy.address}', + timeout=600, + ) as client: + route = f'{cls.VIDEO_TASK_ENDPOINT}/{task_id}' + resp = client.get(route) + try: + data: BytedanceVideoTaskResponse = resp.json() + except Exception: + raise GenerationException + if data.get('id') or data.get('task_id'): + return data + raise GenerationException + + @classmethod + def _wait_video_task_result(cls, task_id: str) -> RunVideoResult: + for _ in range(cls.VIDEO_POLL_ATTEMPTS_LIMIT): + data = cls._retrieve_video_task(task_id=task_id) + status = str(data.get('status', '')).lower() + if status == BytedanceVideoTaskStatus.SUCCEEDED: + return data['content']['video_url'], data['usage']['completion_tokens'] + if status not in { + BytedanceVideoTaskStatus.QUEUED, + BytedanceVideoTaskStatus.RUNNING, + }: + raise GenerationException + time.sleep(cls.VIDEO_POLL_DELAY_SECONDS) + raise GenerationException @@ -21,6 +21,7 @@ from ml_model.services.gemini_3_1 import Gemini_3_1 from ml_model.services.geminiimage import Geminiimage from ml_model.services.gemma import Gemma from ml_model.services.gptimage import Gptimage +from ml_model.services.glm_4_7 import Glm_4_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 @@ -64,6 +65,7 @@ from ml_model.services.reve import Reve from ml_model.services.runway import Runway from ml_model.services.sdxlemoji import Sdxlemoji from ml_model.services.seedance import Seedance +from ml_model.services.seedance_2_dreamina import Seedance_2_Dreamina from ml_model.services.seedream import Seedream from ml_model.services.sora import Sora from ml_model.services.stablediffusion import Stablediffusion @@ -0,0 +1,102 @@ +import time +import logging +from datetime import timedelta +from decimal import Decimal +from typing import Any, Iterator + +from messages.models import Message +from ml_model.adapters.bytedance_model_ark import BytedanceContentType, BytedanceModelArkAdapter +from ml_model.services.base import SimpleService +from ml_model.tasks import bytedance_model_ark_run +from tools.chats.models import Chat +from tools.copywrite.models import Copywrite +from tools.public_api.models import APIStore + +logger = logging.getLogger(__name__) + + +class Glm_4_7(SimpleService): + TOKENS_COST = { + 'glm-4-7-251222': { + 'input': Decimal('300'), + 'output': Decimal('1100'), + }, # 1M tokens + } + + def calculate_price(self, version: str, input_tokens: int, output_tokens: int) -> Decimal: + price_map = self.TOKENS_COST[version] + price = ( + input_tokens * price_map['input'] / 1_000_000 + + output_tokens * price_map['output'] / 1_000_000 + ) + return price.quantize(Decimal('0.1'), rounding='ROUND_UP') + + def save_results(self, content: Iterator[Any], t: timedelta, save: bool = True) -> list[Message]: + msgs = [ + Message( + content=content, + content_object=self.store, + elapsed_time=t, + ) + ] + if save: + return Message.objects.bulk_create(msgs) + return msgs + + def make(self, input_message: Message, save: bool = True) -> list[Message]: + start_time = time.time() + version = input_message.info.pop('version', 'glm-4-7-251222') + callback_data = {**input_message.info} + messages = self.get_chat_history() + messages.append({'role': 'user', 'content': input_message.content}) + if input_message.file: + logger.info('GLM file input ignored: model does not support image/video input') + result = bytedance_model_ark_run( + model=version, + callback_data=callback_data, + content_type=BytedanceContentType.CHAT, + messages=messages, + ) + process_time = timedelta(seconds=(time.time() - start_time)) + self.handle_invoice( + input_message.content_object.model, + version=version, + input_tokens=result[1], + output_tokens=result[2], + ) + msgs = self.save_results(result[0], process_time, save) + return msgs + + def get_chat_history(self, message_limit: int = 10, max_character_limit: int = 1500) -> list[dict[str, str | list]]: + if isinstance(self.store, Chat): + air_messages = list( + reversed( + Message.objects.filter( + chats_chats_messages=self.store, is_deleted=False, is_sent=True + ).order_by('-created_at')[1:message_limit + 1] + ) + ) + elif isinstance(self.store, APIStore): + air_messages = [] + elif isinstance(self.store, Copywrite): + air_messages = list( + reversed( + Message.objects.filter( + copywrite_copywrites_messages=self.store, + is_deleted=False, + is_sent=True, + ).order_by('-created_at')[:message_limit] + ) + ) + memory = [] + for msg in air_messages: + content = msg.content or '' + if msg.from_model: + memory.append({'role': 'assistant', 'content': content}) + else: + memory.append({'role': 'user', 'content': content}) + character_length = sum(len(content['content']) for content in memory) + while character_length > max_character_limit: + character_length -= len(memory.pop(0)['content']) + + return memory @@ -1,4 +1,3 @@ -import base64 import time from datetime import timedelta from decimal import Decimal @@ -0,0 +1,205 @@ +import json +import os +import subprocess +import tempfile +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.utils.translation import gettext as _ +from messages.models import Message +from ml_model.adapters.bytedance_model_ark import BytedanceContentType +from ml_model.exceptions import FileExtensionNotSupported, InvalidParameterError +from ml_model.services.FileService import FileProcessingService +from ml_model.services.base import SimpleService +from ml_model.tasks import bytedance_model_ark_run + +from payments.exceptions.insufficient_balance import InsufficientBalance +from payments.selectors.payment_plan_selector import PaymentPlanSelector + + +class Seedance_2_Dreamina(SimpleService): + TOKENS_COST = { + 'dreamina-seedance-2-0': { + 'non_video_in': {'480p': Decimal('3500'), '720p': Decimal('3500'), '1080p': Decimal('3850')}, + 'video_in': {'480p': Decimal('2150'), '720p': Decimal('2150'), '1080p': Decimal('2350')}, + }, + 'dreamina-seedance-2-0-fast': { + 'non_video_in': {'480p': Decimal('2800'), '720p': Decimal('2800')}, + 'video_in': {'480p': Decimal('1650'), '720p': Decimal('1650')}, + }, + } + + RESOLUTION_RATIO_TO_WIDTH_HEIGHT = { + '480p': { + '16:9': (864, 496), + '4:3': (752, 560), + '1:1': (640, 640), + '3:4': (560, 752), + '9:16': (496, 864), + '21:9': (992, 432), + }, + '720p': { + '16:9': (1280, 720), + '4:3': (1112, 834), + '1:1': (960, 960), + '3:4': (834, 1112), + '9:16': (720, 1280), + '21:9': (1470, 630), + }, + '1080p': { + '16:9': (1920, 1080), + '4:3': (1664, 1248), + '1:1': (1440, 1440), + '3:4': (1248, 1664), + '9:16': (1080, 1920), + '21:9': (2206, 946), + }, + } + + @classmethod + def predict_price(cls, content: str, file_exists: bool, info: dict[str, Any]) -> Decimal | None: + resolution = info['resolution'] + duration = info['duration'] + version = info['version'] + ratio = info.get('ratio', '16:9') + if version == 'dreamina-seedance-2-0-fast' and resolution == '1080p': + return None + generation_type = 'video_in' if file_exists else 'non_video_in' + price = ( + cls.TOKENS_COST[version][generation_type][resolution] + / 1_000_000 + * cls.calculate_estimated_tokens( + resolution=resolution, + ratio=ratio, + duration=duration + (15 if generation_type == 'video_in' else 0), + ) + ) + return price.quantize(Decimal('0.1'), rounding='ROUND_UP') + + def calculate_price( + self, completion_tokens: int, resolution: str, version: str, generation_type: str + ) -> Decimal: + price = self.TOKENS_COST[version][generation_type][resolution] * completion_tokens / 1_000_000 + return price.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]: + version = input_message.info.pop('version', 'dreamina-seedance-2-0') + resolution = input_message.info.get('resolution', '720p') + duration = input_message.info.get('duration', 5) + ratio = input_message.info.get('ratio', '16:9') + if version == 'dreamina-seedance-2-0-fast' and resolution == '1080p': + raise InvalidParameterError(_('1080p output is not supported for Seedance Dreamina 2.0 Fast.')) + file = input_message.file or None + file_extension = None + generation_type = 'non_video_in' + video_length = 0 + if file: + file_bytes = file.read() + kind = filetype.guess(file_bytes[:50]) + file_extension = ( + FileProcessingService.get_file_extension(kind.extension, file_bytes) if kind else None + ) + available_extensions = ('JPG', 'JPEG', 'PNG', 'WEBP', 'MP4') + if not file_extension or file_extension.upper() not in available_extensions: + raise FileExtensionNotSupported(available_extensions) + elif file_extension.upper() == 'MP4': + fd, tmp_path = tempfile.mkstemp(suffix='.mp4') + try: + os.write(fd, file_bytes) + os.close(fd) + fd = -1 + video_length = float(self.duration_seconds(tmp_path)) + finally: + if fd >= 0: + try: + os.close(fd) + except OSError: + pass + try: + os.unlink(tmp_path) + except FileNotFoundError: + pass + generation_type = 'video_in' + predicted = ( + self.TOKENS_COST[version][generation_type][resolution] + / 1_000_000 + * self.calculate_estimated_tokens( + resolution=resolution, ratio=ratio, duration=duration + video_length + ) + ) + 10 + if (balance := PaymentPlanSelector(self.store.user).get_current_balance()) < predicted: + raise InsufficientBalance(balance, predicted) + callback_data = dict({**input_message.info}) + content = [{'type': 'text', 'text': input_message.content}] + if file: + if file_extension and file_extension.upper() == 'MP4': + content.append( + {'type': 'video_url', 'video_url': {'url': file.url}, 'role': 'reference_video'} + ) + else: + content.append( + {'type': 'image_url', 'image_url': {'url': file.url}, 'role': 'reference_image'} + ) + callback_data.update({'content': content}) + start_time = time.time() + video, completion_tokens = bytedance_model_ark_run( + model=f'{version}-260128', + callback_data=callback_data, + content_type=BytedanceContentType.VIDEO, + ) + process_time = timedelta(seconds=(time.time() - start_time)) + self.handle_invoice( + input_message.content_object.model, + completion_tokens=completion_tokens, + resolution=resolution, + version=version, + generation_type=generation_type, + ) + msgs = self.save_results(input_message.content, process_time, video, save) + return msgs + + @classmethod + def calculate_estimated_tokens(cls, resolution: str, ratio: str, duration: int) -> Decimal: + return Decimal( + duration + * cls.RESOLUTION_RATIO_TO_WIDTH_HEIGHT[resolution][ratio][0] + * cls.RESOLUTION_RATIO_TO_WIDTH_HEIGHT[resolution][ratio][1] + * 24 + / 1024 + ) + + @classmethod + def duration_seconds(cls, video_path: str) -> float: + out = subprocess.run( + [ + 'ffprobe', + '-v', + 'error', + '-show_entries', + 'format=duration', + '-of', + 'json', + video_path, + ], + capture_output=True, + text=True, + check=True, + ) + return float(json.loads(out.stdout)["format"]["duration"]) \ No newline at end of file @@ -18,6 +18,7 @@ from deepl.translator import TextResult from requests import Response from backend import settings +from ml_model.adapters.bytedance_model_ark import BytedanceContentType, BytedanceModelArkAdapter from ml_model.exceptions import DeploymentDisabled, ModelTimeoutError, GenerationException, RequestBlocked from ml_model.utils import count_openrouter_tokens from poller.models import Proxy @@ -201,6 +202,23 @@ def upscale_run(payload: dict[str, tuple[str, IO]]) -> list[str]: def evaluate_model(model_name: str, data: Dict[str, Any]): ... +@shared_task +def bytedance_model_ark_run( + model: str, + callback_data: dict, + content_type: str = BytedanceContentType.CHAT, + messages: list | None = None, + include_reasoning: bool = False, +): + return BytedanceModelArkAdapter.run( + model=model, + content_type=BytedanceContentType(content_type), + callback_data=callback_data, + messages=messages, + include_reasoning=include_reasoning, + ) + + @shared_task def drop_redis_vectors(message_uid: str) -> None: redis_client = redis.Redis(host=settings.REDIS_HOST, port=settings.REDIS_PORT, db=0) @@ -18,6 +18,7 @@ DEEPL_API_KEY=4bb58b98-ca95-5978-9be0-ed437df6c15c:fx SERPER_API_KEY=ed8e0dbcc26dacf3f7f99fbc8b3add9ada0c793e FLUX_API_KEY=dccaf377-aecf-4cf0-aff4-dde47cee340d OPENROUTER_API_KEY=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 +BYTEDANCE_MODEL_ARK_API_KEY=ark-151e9e89-7275-4dbf-bbb3-d2b32bb69d61-3ca5b FAL_API_KEY=617f0fe4-c627-4119-9681-11af2c3e416a:3d618ccd0ee11ed82543801d7da96d1d # EXTERNAL SERVICES