| 1 | import logging |
| 2 | from typing import List, Optional |
| 3 | import asyncio |
| 4 | from google import genai |
| 5 | from google.genai import types |
| 6 | from google.genai.errors import ClientError |
| 7 | from interfaces.video_output import VideoOutput |
| 8 | from utils.rate_limiter import RateLimiter |
| 9 | |
| 10 | # https://ai.google.dev/gemini-api/docs/video-generation?hl=zh-cn |
| 11 | |
| 12 | |
| 13 | class VideoGeneratorVeoGoogleAPI: |
| 14 | def __init__( |
| 15 | self, |
| 16 | api_key: str, |
| 17 | t2v_model: str = "veo-3.1-generate-preview", |
| 18 | ff2v_model: str = "veo-3.1-generate-preview", |
| 19 | flf2v_model: str = "veo-3.1-generate-preview", |
| 20 | rate_limiter: Optional[RateLimiter] = None, |
| 21 | ): |
| 22 | self.api_key = api_key |
| 23 | self.t2v_model = t2v_model |
| 24 | self.ff2v_model = ff2v_model |
| 25 | self.flf2v_model = flf2v_model |
| 26 | self.rate_limiter = rate_limiter |
| 27 | |
| 28 | self.client = genai.Client( |
| 29 | api_key=api_key, |
| 30 | ) |
| 31 | |
| 32 | async def generate_single_video( |
| 33 | self, |
| 34 | prompt: str, |
| 35 | reference_image_paths: List[str], |
| 36 | resolution: str = "1080p", |
| 37 | aspect_ratio: str = "16:9", |
| 38 | duration: int = 8, |
| 39 | **kwargs, |
| 40 | ) -> VideoOutput: |
| 41 | |
| 42 | params = { |
| 43 | "prompt": prompt, |
| 44 | } |
| 45 | config_params = { |
| 46 | "resolution": resolution, |
| 47 | "aspect_ratio": aspect_ratio, |
| 48 | "duration_seconds": duration, |
| 49 | } |
| 50 | if len(reference_image_paths) == 0: |
| 51 | params["model"] = self.t2v_model |
| 52 | elif len(reference_image_paths) == 1: |
| 53 | params["model"] = self.ff2v_model |
| 54 | params["image"] = types.Image.from_file(location=reference_image_paths[0]) |
| 55 | elif len(reference_image_paths) == 2: |
| 56 | # First+last-frame ("flf2v") interpolation returns 400 |
| 57 | # INVALID_ARGUMENT ("Your use case is currently not supported") |
| 58 | # on the public Gemini Developer API (it appears to require |
| 59 | # Vertex AI / allowlisting). Fall back to first-frame-only so |
| 60 | # the shot still renders instead of failing the pipeline; the |
| 61 | # clip just is not pinned to the generated last frame. |
| 62 | logging.warning( |
| 63 | "Two reference images provided but first+last-frame video " |
| 64 | "generation is not available on this API key; falling back " |
| 65 | "to first-frame-only generation." |
| 66 | ) |
| 67 | params["model"] = self.ff2v_model |
| 68 | params["image"] = types.Image.from_file(location=reference_image_paths[0]) |
| 69 | else: |
| 70 | raise ValueError("The number of reference images must be no more than 2") |
| 71 | |
| 72 | logging.info(f"Calling {params['model']} to generate video...") |
| 73 | |
| 74 | # Apply rate limiting if configured |
| 75 | if self.rate_limiter: |
| 76 | await self.rate_limiter.acquire() |
| 77 | |
| 78 | # Retry logic for rate limit errors |
| 79 | max_retries = 3 |
| 80 | retry_delay = 5 |
| 81 | |
| 82 | for attempt in range(max_retries): |
| 83 | try: |
| 84 | operation = self.client.models.generate_videos( |
| 85 | **params, |
| 86 | config=types.GenerateVideosConfig(**config_params), |
| 87 | ) |
| 88 | break |
| 89 | except ClientError as e: |
| 90 | # google.genai.errors.ClientError exposes the HTTP status |
| 91 | # as `.code`; `.status_code` does not exist, so this line |
| 92 | # raised AttributeError and masked every real ClientError. |
| 93 | if e.code == 429 and attempt < max_retries - 1: |
| 94 | wait_time = retry_delay * (2 ** attempt) |
| 95 | logging.warning(f"Rate limit hit (429), retrying in {wait_time}s... (attempt {attempt + 1}/{max_retries})") |
| 96 | await asyncio.sleep(wait_time) |
| 97 | else: |
| 98 | raise |
| 99 | |
| 100 | while not operation.done: |
| 101 | await asyncio.sleep(2) |
| 102 | operation = self.client.operations.get(operation) |
| 103 | logging.info(f"Video generation not completed, waiting 2 seconds...") |
| 104 | |
| 105 | # Check if operation completed successfully |
| 106 | if operation.error: |
| 107 | error_msg = f"Video generation failed: {operation.error}" |
| 108 | logging.error(error_msg) |
| 109 | raise RuntimeError(error_msg) |
| 110 | |
| 111 | if not operation.response: |
| 112 | error_msg = "Video generation completed but no response received" |
| 113 | logging.error(error_msg) |
| 114 | raise RuntimeError(error_msg) |
| 115 | |
| 116 | if not hasattr(operation.response, 'generated_videos') or not operation.response.generated_videos: |
| 117 | error_msg = "Video generation completed but no videos were generated" |
| 118 | logging.error(error_msg) |
| 119 | raise RuntimeError(error_msg) |
| 120 | |
| 121 | generated_video = operation.response.generated_videos[0] |
| 122 | self.client.files.download(file=generated_video.video) |
| 123 | |
| 124 | video_output = VideoOutput( |
| 125 | fmt="bytes", |
| 126 | ext="mp4", |
| 127 | data=generated_video.video.video_bytes, |
| 128 | ) |
| 129 | return video_output |
| 130 |