@@ -80,3 +80,4 @@ from ml_model.services.vicuna import Vicuna from ml_model.services.wan import Wan from ml_model.services.wan_lite import Wan_Lite from ml_model.services.whisper import Whisper +from ml_model.services.video_test_model import Video_Test_Model \ No newline at end of file @@ -0,0 +1,130 @@ +import subprocess +import time +from datetime import timedelta +from decimal import Decimal +from io import BytesIO +from typing import Any +import filetype +import json + +import requests +from django.core.files import File + +from messages.models import Message +from ml_model.exceptions import CorruptedFileError, InvalidParameterError +from ml_model.services.base import SimpleService +from payments.exceptions.insufficient_balance import InsufficientBalance +from payments.selectors.payment_plan_selector import PaymentPlanSelector + + + +class Video_Test_Model(SimpleService): + + TOKENS_COST_PER_UNIT = Decimal('15') + TOKENS_COST_PER_SECOND = Decimal('0.5') + + PLACEHOLDER_URL='https://imgur.com/QPLhtj1.mp4' + + def calculate_price(self, cps: str = 'per-unit', duration: float = 1, num_videos: int = 1) -> Decimal: + if cps == 'per-second': + return (self.TOKENS_COST_PER_SECOND * Decimal(duration) * num_videos).quantize( + Decimal('0.1'), rounding='ROUND_UP' + ) + return self.TOKENS_COST_PER_UNIT * num_videos + + @classmethod + def predict_price(cls, content: str, file_exists: bool, info: dict[str, Any]) -> Decimal | None: + num_videos = info.get('num_videos', 1) + + return cls.TOKENS_COST_PER_UNIT * num_videos + + def save_results( + self, + content: str, + t: timedelta, + videos: list[bytes], + save: bool = True, + ) -> list[Message]: + messages: list[Message] = [] + for video in videos: + messages.append( + Message( + content_object=self.store, + elapsed_time=t, + content=content, + file=File(BytesIO(video), '.mp4'), + ) + ) + if save: + return Message.objects.bulk_create(messages) + return messages + + + def make(self, input_message: Message, save: bool = True) -> list[Message]: + cps = input_message.info.get('cps', 'per-unit') + cvu = input_message.info.get('cvu') or self.PLACEHOLDER_URL + num_videos = input_message.info.get('num_videos', 1) + + start_time = time.time() + + video_bytes = self._fetch_video(cvu) + duration=self._get_duration(video_bytes) + + if (balance := PaymentPlanSelector(self.store.user).get_current_balance()) < (cost := self.calculate_price(cps, duration, num_videos)): + raise InsufficientBalance(balance, cost) + + + videos = [video_bytes] * num_videos + + process_time = timedelta(seconds=(time.time() - start_time)) + self.handle_invoice(input_message.content_object.model, cps, duration, num_videos) + + msgs = self.save_results(input_message.content, process_time, videos, save) + return msgs + + + + def _fetch_video(self, url: str): + headers = { + "User-Agent": ( + "Mozilla/5.0 (Windows NT 10.0; Win64; x64) " + "AppleWebKit/537.36 (KHTML, like Gecko) " + "Chrome/137.0 Safari/537.36" + ) + } + try: + response = requests.get( + url, + headers=headers, + timeout=10 + ) + response.raise_for_status() + except requests.RequestException as e: + raise InvalidParameterError(f'Invalid video URL: {e}') + + kind = filetype.guess(response.content[:120]) + + if not kind: + raise CorruptedFileError + + if not kind.mime.startswith('video/'): + raise InvalidParameterError('Video format not supported') + + return response.content + + + def _get_duration(self, video_bytes: bytes) -> float: + result = subprocess.run( + [ + "ffprobe", + "-v", "quiet", + "-print_format", "json", + "-show_format", + "-", + ], + input=video_bytes, + capture_output=True, + ) + + data = json.loads(result.stdout) + return float(data["format"]["duration"]) \ No newline at end of file