| 1 | import logging |
| 2 | from typing import List, Optional |
| 3 | from PIL import Image |
| 4 | import asyncio |
| 5 | import aiohttp |
| 6 | import os |
| 7 | from interfaces.video_output import VideoOutput |
| 8 | from utils.image import image_path_to_b64 |
| 9 | |
| 10 | |
| 11 | def _env_int(name: str, default: int) -> int: |
| 12 | try: |
| 13 | return max(0, int(os.environ.get(name, str(default)))) |
| 14 | except ValueError: |
| 15 | return default |
| 16 | |
| 17 | |
| 18 | def _env_float(name: str, default: float) -> float: |
| 19 | try: |
| 20 | return max(0.0, float(os.environ.get(name, str(default)))) |
| 21 | except ValueError: |
| 22 | return default |
| 23 | |
| 24 | |
| 25 | def _emit_progress(progress, stage: str, message: str, metadata: dict | None = None) -> None: |
| 26 | if progress is not None: |
| 27 | progress(stage, message, metadata or {}) |
| 28 | |
| 29 | |
| 30 | class VideoGeneratorVeoYunwuAPI: |
| 31 | def __init__( |
| 32 | self, |
| 33 | api_key: str, |
| 34 | t2v_model: str = "veo3.1-fast", # text to video |
| 35 | ff2v_model: str = "veo3.1-fast", # first frame to video |
| 36 | flf2v_model: str = "veo2-fast-frames", # first and last frame to video |
| 37 | base_url: str = "https://yunwu.ai", |
| 38 | ): |
| 39 | """ |
| 40 | all models: |
| 41 | veo2 |
| 42 | veo2-fast |
| 43 | veo2-fast-frames |
| 44 | veo2-fast-components |
| 45 | veo2-pro |
| 46 | veo3 |
| 47 | veo3-fast |
| 48 | veo3-pro |
| 49 | veo3-pro-frames |
| 50 | veo3-fast-frames |
| 51 | veo3-frames |
| 52 | |
| 53 | NOTE: veo3 does not support first and last frame to video generation. |
| 54 | """ |
| 55 | self.base_url = base_url.rstrip("/") |
| 56 | self.api_key = api_key |
| 57 | self.t2v_model = t2v_model |
| 58 | self.ff2v_model = ff2v_model |
| 59 | self.flf2v_model = flf2v_model |
| 60 | |
| 61 | async def generate_single_video( |
| 62 | self, |
| 63 | prompt: str = "", |
| 64 | reference_image_paths: List[Image.Image] = [], |
| 65 | aspect_ratio: str = "16:9", |
| 66 | **kwargs, |
| 67 | ) -> VideoOutput: |
| 68 | progress = kwargs.get("progress") |
| 69 | create_retries = _env_int("VIMAX_VIDEO_CREATE_RETRIES", 3) |
| 70 | query_timeout_seconds = _env_float("VIMAX_VIDEO_QUERY_TIMEOUT_SECONDS", 600.0) |
| 71 | request_timeout_seconds = _env_float("VIMAX_VIDEO_REQUEST_TIMEOUT_SECONDS", 60.0) |
| 72 | poll_interval_seconds = _env_float("VIMAX_VIDEO_POLL_INTERVAL_SECONDS", 5.0) |
| 73 | max_query_errors = _env_int("VIMAX_VIDEO_MAX_QUERY_ERRORS", 5) |
| 74 | if len(reference_image_paths) == 0: |
| 75 | model = self.t2v_model |
| 76 | elif len(reference_image_paths) == 1: |
| 77 | model = self.ff2v_model |
| 78 | elif len(reference_image_paths) == 2: |
| 79 | model = self.flf2v_model |
| 80 | else: |
| 81 | raise ValueError("The number of reference images must be no more than 2") |
| 82 | |
| 83 | logging.info(f"Calling {model} to generate video...") |
| 84 | |
| 85 | # 1. Create video generation task |
| 86 | payload = { |
| 87 | "prompt": prompt, |
| 88 | "model": model, |
| 89 | "images": [image_path_to_b64(image_path, mime=True) for image_path in reference_image_paths], |
| 90 | "enhance_prompt": True, |
| 91 | } |
| 92 | # only veo3 supports aspect ratio setting |
| 93 | if model.startswith("veo3"): |
| 94 | payload["aspect_ratio"] = aspect_ratio |
| 95 | |
| 96 | headers = { |
| 97 | "Accept": "application/json", |
| 98 | "Authorization": f"Bearer {self.api_key}", |
| 99 | "Content-Type": "application/json", |
| 100 | } |
| 101 | |
| 102 | url = f"{self.base_url}/v1/video/create" |
| 103 | task_id = None |
| 104 | last_create_error = None |
| 105 | timeout = aiohttp.ClientTimeout(total=request_timeout_seconds) |
| 106 | for attempt in range(1, create_retries + 1): |
| 107 | try: |
| 108 | _emit_progress(progress, "video_create", f"Creating video generation task with {model}", {"model": model, "attempt": attempt, "max_attempts": create_retries}) |
| 109 | async with aiohttp.ClientSession(timeout=timeout) as session: |
| 110 | async with session.post(url, headers=headers, json=payload) as response: |
| 111 | response_payload = await response.json(content_type=None) |
| 112 | logging.debug(f"Response: {response_payload}") |
| 113 | if response.status >= 400: |
| 114 | raise RuntimeError(f"Video create failed with HTTP {response.status}: {response_payload}") |
| 115 | task_id = response_payload.get("id") |
| 116 | if not task_id: |
| 117 | raise RuntimeError(f"Video create response missing id: {response_payload}") |
| 118 | logging.info(f"Video generation task created successfully. Task ID: {task_id}") |
| 119 | _emit_progress(progress, "video_task_created", "Video generation task created", {"model": model, "task_id": task_id}) |
| 120 | break |
| 121 | except Exception as e: |
| 122 | last_create_error = e |
| 123 | logging.error(f"Error occurred while creating video generation task: {e}.") |
| 124 | _emit_progress(progress, "video_create_error", f"Video create attempt {attempt} failed", {"model": model, "attempt": attempt, "error": str(e)}) |
| 125 | if attempt < create_retries: |
| 126 | await asyncio.sleep(1) |
| 127 | if not task_id: |
| 128 | raise RuntimeError(f"Video create failed after {create_retries} attempts: {last_create_error}") |
| 129 | |
| 130 | |
| 131 | # 2. Query the video generation task until the video generation is completed |
| 132 | headers = { |
| 133 | 'Accept': 'application/json', |
| 134 | 'Authorization': f'Bearer {self.api_key}', |
| 135 | } |
| 136 | |
| 137 | deadline = asyncio.get_running_loop().time() + query_timeout_seconds if query_timeout_seconds > 0 else None |
| 138 | query_errors = 0 |
| 139 | last_status = None |
| 140 | while deadline is None or asyncio.get_running_loop().time() < deadline: |
| 141 | try: |
| 142 | async with aiohttp.ClientSession(timeout=timeout) as session: |
| 143 | async with session.get(f"{self.base_url}/v1/video/query?id={task_id}", headers=headers) as response: |
| 144 | payload = await response.json(content_type=None) |
| 145 | logging.debug(f"Response: {payload}") |
| 146 | if response.status >= 400: |
| 147 | raise RuntimeError(f"Video query failed with HTTP {response.status}: {payload}") |
| 148 | status = payload.get("status") |
| 149 | if not status: |
| 150 | raise RuntimeError(f"Video query response missing status: {payload}") |
| 151 | query_errors = 0 |
| 152 | except Exception as e: |
| 153 | query_errors += 1 |
| 154 | logging.error(f"Error occurred while querying video generation task: {e}.") |
| 155 | _emit_progress(progress, "video_query_error", "Video query failed", {"model": model, "task_id": task_id, "error": str(e), "query_errors": query_errors, "max_query_errors": max_query_errors}) |
| 156 | if query_errors >= max_query_errors: |
| 157 | raise RuntimeError(f"Video query failed {query_errors} times for task {task_id}: {e}") |
| 158 | await asyncio.sleep(poll_interval_seconds) |
| 159 | continue |
| 160 | |
| 161 | if status == "completed": |
| 162 | logging.info(f"Video generation completed successfully") |
| 163 | video_url = payload.get("video_url") |
| 164 | if not video_url: |
| 165 | raise RuntimeError(f"Video task completed without video_url: {payload}") |
| 166 | _emit_progress(progress, "video_completed", "Video generation completed", {"model": model, "task_id": task_id}) |
| 167 | return VideoOutput(fmt="url", ext="mp4", data=video_url) |
| 168 | elif status == "failed": |
| 169 | logging.error(f"Video generation failed: \n{payload}") |
| 170 | raise RuntimeError(f"Video generation failed for task {task_id}: {payload}") |
| 171 | else: |
| 172 | logging.info(f"Video generation status: {status}, waiting 1 second...") |
| 173 | last_status = status |
| 174 | _emit_progress(progress, "video_status", f"Video generation status: {status}", {"model": model, "task_id": task_id, "status": status}) |
| 175 | await asyncio.sleep(poll_interval_seconds) |
| 176 | continue |
| 177 | raise RuntimeError(f"Video generation timed out after {query_timeout_seconds:g}s for task {task_id}; last_status={last_status}") |
| 178 |