| 1 | import asyncio |
| 2 | import logging |
| 3 | import os |
| 4 | from typing import List |
| 5 | from urllib.parse import urljoin |
| 6 | |
| 7 | import aiohttp |
| 8 | |
| 9 | from interfaces.video_output import VideoOutput |
| 10 | from utils.image import image_path_to_b64 |
| 11 | |
| 12 | |
| 13 | def _env_int(name: str, default: int) -> int: |
| 14 | try: |
| 15 | return max(0, int(os.environ.get(name, str(default)))) |
| 16 | except ValueError: |
| 17 | return default |
| 18 | |
| 19 | |
| 20 | def _env_float(name: str, default: float) -> float: |
| 21 | try: |
| 22 | return max(0.0, float(os.environ.get(name, str(default)))) |
| 23 | except ValueError: |
| 24 | return default |
| 25 | |
| 26 | |
| 27 | def _env_bool(name: str, default: bool) -> bool: |
| 28 | raw = os.environ.get(name) |
| 29 | if raw is None: |
| 30 | return default |
| 31 | return raw.strip().lower() in {"1", "true", "yes", "on"} |
| 32 | |
| 33 | |
| 34 | def _emit_progress(progress, stage: str, message: str, metadata: dict | None = None) -> None: |
| 35 | if progress is not None: |
| 36 | progress(stage, message, metadata or {}) |
| 37 | |
| 38 | |
| 39 | class VideoGeneratorOpenRouterAPI: |
| 40 | def __init__( |
| 41 | self, |
| 42 | api_key: str, |
| 43 | model: str = "google/veo-3.1-lite", |
| 44 | base_url: str = "https://openrouter.ai/api/v1", |
| 45 | http_referer: str = "", |
| 46 | app_title: str = "ViMax", |
| 47 | ): |
| 48 | self.api_key = api_key |
| 49 | self.model = model |
| 50 | self.base_url = base_url.rstrip("/") |
| 51 | self.http_referer = http_referer |
| 52 | self.app_title = app_title |
| 53 | |
| 54 | async def generate_single_video( |
| 55 | self, |
| 56 | prompt: str = "", |
| 57 | reference_image_paths: List[str] = [], |
| 58 | aspect_ratio: str = "16:9", |
| 59 | **kwargs, |
| 60 | ) -> VideoOutput: |
| 61 | progress = kwargs.get("progress") |
| 62 | request_timeout_seconds = _env_float("VIMAX_VIDEO_REQUEST_TIMEOUT_SECONDS", 60.0) |
| 63 | query_timeout_seconds = _env_float("VIMAX_VIDEO_QUERY_TIMEOUT_SECONDS", 600.0) |
| 64 | poll_interval_seconds = _env_float("VIMAX_VIDEO_POLL_INTERVAL_SECONDS", 10.0) |
| 65 | duration = _env_int("VIMAX_OPENROUTER_VIDEO_DURATION", 8) |
| 66 | resolution = os.environ.get("VIMAX_OPENROUTER_VIDEO_RESOLUTION", "720p") |
| 67 | generate_audio = _env_bool("VIMAX_OPENROUTER_GENERATE_AUDIO", True) |
| 68 | |
| 69 | payload = { |
| 70 | "model": self.model, |
| 71 | "prompt": prompt, |
| 72 | "aspect_ratio": aspect_ratio, |
| 73 | "duration": duration, |
| 74 | "resolution": resolution, |
| 75 | "generate_audio": generate_audio, |
| 76 | } |
| 77 | frame_images = _frame_images(reference_image_paths) |
| 78 | if frame_images: |
| 79 | payload["frame_images"] = frame_images |
| 80 | |
| 81 | headers = self._headers() |
| 82 | timeout = aiohttp.ClientTimeout(total=request_timeout_seconds) |
| 83 | _emit_progress(progress, "video_create", f"Creating OpenRouter video generation task with {self.model}", {"model": self.model, "duration": duration, "resolution": resolution, "frame_count": len(frame_images)}) |
| 84 | |
| 85 | create_status, create_payload = await _post_json( |
| 86 | f"{self.base_url}/videos", |
| 87 | headers=headers, |
| 88 | payload=payload, |
| 89 | timeout=timeout, |
| 90 | hard_timeout_seconds=request_timeout_seconds, |
| 91 | ) |
| 92 | if create_status >= 400: |
| 93 | raise RuntimeError(f"OpenRouter video create failed with HTTP {create_status}: {create_payload}") |
| 94 | job_id = create_payload.get("id") |
| 95 | polling_url = create_payload.get("polling_url") |
| 96 | if not job_id or not polling_url: |
| 97 | raise RuntimeError(f"OpenRouter video create response missing id or polling_url: {create_payload}") |
| 98 | _emit_progress(progress, "video_task_created", "OpenRouter video generation task created", {"model": self.model, "job_id": job_id, "status": create_payload.get("status")}) |
| 99 | |
| 100 | poll_url = _absolute_url(self.base_url, polling_url) |
| 101 | deadline = asyncio.get_running_loop().time() + query_timeout_seconds if query_timeout_seconds > 0 else None |
| 102 | last_status = create_payload.get("status") |
| 103 | last_payload = create_payload |
| 104 | while deadline is None or asyncio.get_running_loop().time() < deadline: |
| 105 | await asyncio.sleep(poll_interval_seconds) |
| 106 | poll_status, poll_payload = await _get_json( |
| 107 | poll_url, |
| 108 | headers=headers, |
| 109 | timeout=timeout, |
| 110 | hard_timeout_seconds=request_timeout_seconds, |
| 111 | ) |
| 112 | if poll_status >= 400: |
| 113 | raise RuntimeError(f"OpenRouter video poll failed with HTTP {poll_status}: {poll_payload}") |
| 114 | last_payload = poll_payload |
| 115 | status = poll_payload.get("status") |
| 116 | last_status = status |
| 117 | _emit_progress(progress, "video_status", f"OpenRouter video generation status: {status}", {"model": self.model, "job_id": job_id, "status": status}) |
| 118 | |
| 119 | if status == "completed": |
| 120 | urls = poll_payload.get("unsigned_urls") or [] |
| 121 | if urls: |
| 122 | content_url = urls[0] |
| 123 | else: |
| 124 | content_url = f"{self.base_url}/videos/{job_id}/content?index=0" |
| 125 | _emit_progress(progress, "video_download_start", "Downloading OpenRouter video output", {"model": self.model, "job_id": job_id}) |
| 126 | download_status, data = await _get_bytes( |
| 127 | content_url, |
| 128 | headers=headers if _needs_authorization(content_url) else {}, |
| 129 | timeout=timeout, |
| 130 | hard_timeout_seconds=request_timeout_seconds, |
| 131 | ) |
| 132 | if download_status >= 400: |
| 133 | raise RuntimeError(f"OpenRouter video content download failed with HTTP {download_status}: {data[:500]!r}") |
| 134 | _emit_progress(progress, "video_completed", "OpenRouter video generation completed and downloaded", {"model": self.model, "job_id": job_id}) |
| 135 | return VideoOutput(fmt="bytes", ext="mp4", data=data) |
| 136 | if status in {"failed", "cancelled", "expired"}: |
| 137 | raise RuntimeError(f"OpenRouter video generation {status} for job {job_id}: {poll_payload.get('error') or poll_payload}") |
| 138 | |
| 139 | raise RuntimeError(f"OpenRouter video generation timed out after {query_timeout_seconds:g}s for job {job_id}; last_status={last_status}; last_payload={last_payload}") |
| 140 | |
| 141 | def _headers(self) -> dict[str, str]: |
| 142 | headers = { |
| 143 | "Authorization": f"Bearer {self.api_key}", |
| 144 | "Content-Type": "application/json", |
| 145 | } |
| 146 | if self.http_referer: |
| 147 | headers["HTTP-Referer"] = self.http_referer |
| 148 | if self.app_title: |
| 149 | headers["X-OpenRouter-Title"] = self.app_title |
| 150 | return headers |
| 151 | |
| 152 | |
| 153 | def _frame_images(reference_image_paths: List[str]) -> list[dict]: |
| 154 | if len(reference_image_paths) > 2: |
| 155 | raise ValueError("OpenRouter video generation supports at most first and last frame images") |
| 156 | frame_types = ["first_frame", "last_frame"] |
| 157 | return [ |
| 158 | { |
| 159 | "type": "image_url", |
| 160 | "image_url": {"url": image_path_to_b64(path, mime=True)}, |
| 161 | "frame_type": frame_types[index], |
| 162 | } |
| 163 | for index, path in enumerate(reference_image_paths) |
| 164 | ] |
| 165 | |
| 166 | |
| 167 | def _absolute_url(base_url: str, url: str) -> str: |
| 168 | if url.startswith("http://") or url.startswith("https://"): |
| 169 | return url |
| 170 | return urljoin(f"{base_url.rstrip('/')}/", url.lstrip("/")) |
| 171 | |
| 172 | |
| 173 | def _needs_authorization(url: str) -> bool: |
| 174 | return url.startswith("https://openrouter.ai/api/") |
| 175 | |
| 176 | |
| 177 | async def _post_json(url: str, *, headers: dict[str, str], payload: dict, timeout: aiohttp.ClientTimeout, hard_timeout_seconds: float) -> tuple[int, dict]: |
| 178 | async def request() -> tuple[int, dict]: |
| 179 | async with aiohttp.ClientSession(timeout=timeout) as session: |
| 180 | async with session.post(url, headers=headers, json=payload) as response: |
| 181 | return response.status, await response.json(content_type=None) |
| 182 | |
| 183 | return await asyncio.wait_for(request(), timeout=hard_timeout_seconds + 5) |
| 184 | |
| 185 | |
| 186 | async def _get_json(url: str, *, headers: dict[str, str], timeout: aiohttp.ClientTimeout, hard_timeout_seconds: float) -> tuple[int, dict]: |
| 187 | async def request() -> tuple[int, dict]: |
| 188 | async with aiohttp.ClientSession(timeout=timeout) as session: |
| 189 | async with session.get(url, headers=headers) as response: |
| 190 | return response.status, await response.json(content_type=None) |
| 191 | |
| 192 | return await asyncio.wait_for(request(), timeout=hard_timeout_seconds + 5) |
| 193 | |
| 194 | |
| 195 | async def _get_bytes(url: str, *, headers: dict[str, str], timeout: aiohttp.ClientTimeout, hard_timeout_seconds: float) -> tuple[int, bytes]: |
| 196 | async def request() -> tuple[int, bytes]: |
| 197 | async with aiohttp.ClientSession(timeout=timeout) as session: |
| 198 | async with session.get(url, headers=headers) as response: |
| 199 | return response.status, await response.read() |
| 200 | |
| 201 | return await asyncio.wait_for(request(), timeout=hard_timeout_seconds + 5) |
| 202 |