| 1 | import asyncio |
| 2 | import logging |
| 3 | from typing import List, Optional |
| 4 | |
| 5 | import aiohttp |
| 6 | |
| 7 | from interfaces.video_output import VideoOutput |
| 8 | from utils.image import image_path_to_b64 |
| 9 | from utils.rate_limiter import RateLimiter |
| 10 | |
| 11 | |
| 12 | class VideoGeneratorOmniYunwuAPI: |
| 13 | def __init__( |
| 14 | self, |
| 15 | api_key: str, |
| 16 | t2v_model: str = "omni-flash", |
| 17 | i2v_model: str = "omni-flash", |
| 18 | base_url: str = "https://yunwu.ai", |
| 19 | seconds: int = 8, |
| 20 | enable_upsample: bool = False, |
| 21 | enable_sample: Optional[bool] = None, |
| 22 | poll_interval: int = 2, |
| 23 | max_poll_attempts: Optional[int] = 300, |
| 24 | max_create_attempts: int = 3, |
| 25 | rate_limiter: Optional[RateLimiter] = None, |
| 26 | ): |
| 27 | self.api_key = api_key |
| 28 | self.t2v_model = t2v_model |
| 29 | self.i2v_model = i2v_model |
| 30 | self.base_url = base_url.rstrip("/") |
| 31 | self.seconds = seconds |
| 32 | self.enable_upsample = enable_upsample |
| 33 | self.enable_sample = enable_sample |
| 34 | self.poll_interval = poll_interval |
| 35 | self.max_poll_attempts = max_poll_attempts |
| 36 | self.max_create_attempts = max_create_attempts |
| 37 | self.rate_limiter = rate_limiter |
| 38 | |
| 39 | def _headers(self) -> dict: |
| 40 | return { |
| 41 | "Accept": "application/json", |
| 42 | "Authorization": f"Bearer {self.api_key}", |
| 43 | "Content-Type": "application/json", |
| 44 | } |
| 45 | |
| 46 | def _image_uri(self, image_path: str) -> str: |
| 47 | if image_path.startswith(("http://", "https://", "data:")): |
| 48 | return image_path |
| 49 | return image_path_to_b64(image_path, mime=True) |
| 50 | |
| 51 | def _build_payload( |
| 52 | self, |
| 53 | prompt: str, |
| 54 | reference_image_paths: List[str], |
| 55 | aspect_ratio: str, |
| 56 | seconds: Optional[int], |
| 57 | size: Optional[str], |
| 58 | enable_upsample: Optional[bool], |
| 59 | enable_sample: Optional[bool], |
| 60 | ) -> dict: |
| 61 | if len(reference_image_paths) > 3: |
| 62 | raise ValueError("The number of reference images must be no more than 3") |
| 63 | |
| 64 | payload = { |
| 65 | "model": self.t2v_model if len(reference_image_paths) == 0 else self.i2v_model, |
| 66 | "prompt": prompt, |
| 67 | "seconds": str(seconds or self.seconds), |
| 68 | } |
| 69 | |
| 70 | if len(reference_image_paths) == 0: |
| 71 | payload["type"] = 1 |
| 72 | elif len(reference_image_paths) <= 2: |
| 73 | payload["type"] = 2 |
| 74 | payload["images"] = [self._image_uri(path) for path in reference_image_paths] |
| 75 | else: |
| 76 | payload["type"] = 3 |
| 77 | payload["images"] = [self._image_uri(path) for path in reference_image_paths] |
| 78 | |
| 79 | if aspect_ratio: |
| 80 | payload["aspect_ratio"] = aspect_ratio |
| 81 | if size: |
| 82 | payload["size"] = size |
| 83 | if enable_upsample is not None: |
| 84 | payload["enable_upsample"] = enable_upsample |
| 85 | if enable_sample is not None: |
| 86 | payload["enable_sample"] = enable_sample |
| 87 | |
| 88 | return payload |
| 89 | |
| 90 | async def create_video_generation_task( |
| 91 | self, |
| 92 | prompt: str, |
| 93 | reference_image_paths: List[str], |
| 94 | aspect_ratio: str = "16:9", |
| 95 | seconds: Optional[int] = None, |
| 96 | size: Optional[str] = None, |
| 97 | enable_upsample: Optional[bool] = None, |
| 98 | enable_sample: Optional[bool] = None, |
| 99 | ) -> tuple[str, str]: |
| 100 | payload = self._build_payload( |
| 101 | prompt=prompt, |
| 102 | reference_image_paths=reference_image_paths, |
| 103 | aspect_ratio=aspect_ratio, |
| 104 | seconds=seconds, |
| 105 | size=size, |
| 106 | enable_upsample=self.enable_upsample if enable_upsample is None else enable_upsample, |
| 107 | enable_sample=self.enable_sample if enable_sample is None else enable_sample, |
| 108 | ) |
| 109 | |
| 110 | logging.info("Calling %s to generate video...", payload["model"]) |
| 111 | |
| 112 | if self.rate_limiter: |
| 113 | await self.rate_limiter.acquire() |
| 114 | |
| 115 | url = f"{self.base_url}/v1/video/create" |
| 116 | last_error = None |
| 117 | for attempt in range(1, self.max_create_attempts + 1): |
| 118 | try: |
| 119 | async with aiohttp.ClientSession() as session: |
| 120 | async with session.post(url, headers=self._headers(), json=payload) as response: |
| 121 | response_json = await response.json() |
| 122 | http_status = response.status |
| 123 | logging.debug("Response: %s", response_json) |
| 124 | except Exception as e: |
| 125 | last_error = e |
| 126 | logging.error( |
| 127 | "Error occurred while creating video generation task (attempt %s/%s): %s", |
| 128 | attempt, |
| 129 | self.max_create_attempts, |
| 130 | e, |
| 131 | ) |
| 132 | if attempt < self.max_create_attempts: |
| 133 | await asyncio.sleep(attempt) |
| 134 | continue |
| 135 | |
| 136 | if http_status >= 400: |
| 137 | message = f"Video generation task creation failed with HTTP {http_status}: {response_json}" |
| 138 | if http_status < 500: |
| 139 | raise RuntimeError(message) |
| 140 | last_error = RuntimeError(message) |
| 141 | logging.error("%s (attempt %s/%s)", message, attempt, self.max_create_attempts) |
| 142 | if attempt < self.max_create_attempts: |
| 143 | await asyncio.sleep(attempt) |
| 144 | continue |
| 145 | |
| 146 | task_id = response_json.get("id") |
| 147 | if not task_id: |
| 148 | raise RuntimeError(f"Video generation task creation returned no task id: {response_json}") |
| 149 | logging.info("Video generation task created successfully. Task ID: %s", task_id) |
| 150 | return task_id, payload["model"] |
| 151 | |
| 152 | raise RuntimeError( |
| 153 | f"Failed to create video generation task after {self.max_create_attempts} attempts." |
| 154 | ) from last_error |
| 155 | |
| 156 | async def query_video_generation_task(self, task_id: str, model: str) -> str: |
| 157 | url = f"{self.base_url}/v1/video/query" |
| 158 | params = {"id": task_id, "model": model} |
| 159 | |
| 160 | attempts = 0 |
| 161 | while True: |
| 162 | if self.max_poll_attempts is not None and attempts >= self.max_poll_attempts: |
| 163 | raise TimeoutError(f"Video generation did not complete after {attempts} polls.") |
| 164 | attempts += 1 |
| 165 | |
| 166 | try: |
| 167 | async with aiohttp.ClientSession() as session: |
| 168 | async with session.get(url, headers=self._headers(), params=params) as response: |
| 169 | response_json = await response.json() |
| 170 | logging.debug("Response: %s", response_json) |
| 171 | except Exception as e: |
| 172 | logging.error( |
| 173 | "Error occurred while querying video generation task: %s. Retrying in %s seconds...", |
| 174 | e, |
| 175 | self.poll_interval, |
| 176 | ) |
| 177 | await asyncio.sleep(self.poll_interval) |
| 178 | continue |
| 179 | |
| 180 | status = response_json.get("status") |
| 181 | if status == "completed": |
| 182 | detail = response_json.get("detail") or {} |
| 183 | video_url = ( |
| 184 | response_json.get("video_url") |
| 185 | or detail.get("upsample_video_url") |
| 186 | or detail.get("video_url") |
| 187 | ) |
| 188 | if not video_url: |
| 189 | raise RuntimeError(f"Video generation completed without a video URL: {response_json}") |
| 190 | logging.info("Video generation completed successfully. Video URL: %s", video_url) |
| 191 | return video_url |
| 192 | |
| 193 | if status in {"failed", "error"}: |
| 194 | raise RuntimeError(f"Video generation failed: {response_json}") |
| 195 | |
| 196 | logging.info("Video generation status: %s, waiting %s seconds...", status, self.poll_interval) |
| 197 | await asyncio.sleep(self.poll_interval) |
| 198 | |
| 199 | async def generate_single_video( |
| 200 | self, |
| 201 | prompt: str, |
| 202 | reference_image_paths: List[str], |
| 203 | aspect_ratio: str = "16:9", |
| 204 | seconds: Optional[int] = None, |
| 205 | size: Optional[str] = None, |
| 206 | enable_upsample: Optional[bool] = None, |
| 207 | enable_sample: Optional[bool] = None, |
| 208 | **kwargs, |
| 209 | ) -> VideoOutput: |
| 210 | task_id, model = await self.create_video_generation_task( |
| 211 | prompt=prompt, |
| 212 | reference_image_paths=reference_image_paths, |
| 213 | aspect_ratio=aspect_ratio, |
| 214 | seconds=seconds, |
| 215 | size=size, |
| 216 | enable_upsample=enable_upsample, |
| 217 | enable_sample=enable_sample, |
| 218 | ) |
| 219 | video_url = await self.query_video_generation_task(task_id, model) |
| 220 | return VideoOutput(fmt="url", ext="mp4", data=video_url) |
| 221 | |
| 222 | |
| 223 | class VideoGeneratorOminiYunwuAPI(VideoGeneratorOmniYunwuAPI): |
| 224 | """Backward-compatible alias for the common "omini" spelling.""" |
| 225 |