@@ -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,140 @@ +import math +import subprocess +import time +from datetime import timedelta +from decimal import Decimal +from io import BytesIO +from typing import Any, Literal +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): + + RATES = { + 'per-unit': Decimal('15'), + 'per-second': Decimal('0.5'), + } + + + PLACEHOLDER_URL='https://imgur.com/QPLhtj1.mp4' + + def calculate_price(self, strategy: Literal['per-unit','per-second'], duration: int = 1, num_videos: int = 1) -> Decimal: + rate = self.RATES[strategy] + return (rate * duration * num_videos).quantize(Decimal('0.1'), rounding='ROUND_UP') + + + @classmethod + def predict_price(cls, content: str, file_exists: bool, info: dict[str, Any]) -> Decimal | None: + num_videos = info.get('num_videos', 1) + cps = info.get('cps', 'per-unit') + + if cps == 'per-second': + return None + + return cls.RATES['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) + + if cps == 'per-second': + duration = self._get_duration(video_bytes) + else: + duration = 1 + + 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=600 + ) + response.raise_for_status() + except requests.RequestException as exc: + raise InvalidParameterError(f'Invalid video URL: {exc}') + + 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) -> int: + result = subprocess.run( + [ + "ffprobe", + "-v", "quiet", + "-print_format", "json", + "-show_format", + "-", + ], + input=video_bytes, + capture_output=True, + ) + + data = json.loads(result.stdout) + return math.ceil(float(data["format"]["duration"])) \ No newline at end of file