chore: import upstream snapshot with attribution
This commit is contained in:
wehub-resource-sync
@@ -0,0 +1,35 @@
|
||||
# rendering abstraction
|
||||
from .protocols import ImageGenerator, VideoGenerator
|
||||
from .render_backend import RenderBackend
|
||||
|
||||
# image generators
|
||||
from .image_generator_doubao_seedream_yunwu_api import ImageGeneratorDoubaoSeedreamYunwuAPI
|
||||
from .image_generator_nanobanana_google_api import ImageGeneratorNanobananaGoogleAPI
|
||||
from .image_generator_nanobanana_yunwu_api import ImageGeneratorNanobananaYunwuAPI
|
||||
|
||||
# reranker for rag
|
||||
from .reranker_bge_silicon_api import RerankerBgeSiliconapi
|
||||
|
||||
# video generators
|
||||
from .video_generator_doubao_seedance_yunwu_api import VideoGeneratorDoubaoSeedanceYunwuAPI
|
||||
from .video_generator_omni_yunwu_api import VideoGeneratorOmniYunwuAPI, VideoGeneratorOminiYunwuAPI
|
||||
from .video_generator_openrouter_api import VideoGeneratorOpenRouterAPI
|
||||
from .video_generator_veo_google_api import VideoGeneratorVeoGoogleAPI
|
||||
from .video_generator_veo_yunwu_api import VideoGeneratorVeoYunwuAPI
|
||||
|
||||
|
||||
__all__ = [
|
||||
"ImageGenerator",
|
||||
"VideoGenerator",
|
||||
"RenderBackend",
|
||||
"ImageGeneratorDoubaoSeedreamYunwuAPI",
|
||||
"ImageGeneratorNanobananaGoogleAPI",
|
||||
"ImageGeneratorNanobananaYunwuAPI",
|
||||
"RerankerBgeSiliconapi",
|
||||
"VideoGeneratorDoubaoSeedanceYunwuAPI",
|
||||
"VideoGeneratorOmniYunwuAPI",
|
||||
"VideoGeneratorOminiYunwuAPI",
|
||||
"VideoGeneratorOpenRouterAPI",
|
||||
"VideoGeneratorVeoGoogleAPI",
|
||||
"VideoGeneratorVeoYunwuAPI",
|
||||
]
|
||||
@@ -0,0 +1,74 @@
|
||||
# https://yunwu.apifox.cn/api-347960869
|
||||
|
||||
import asyncio
|
||||
import logging
|
||||
import aiohttp
|
||||
from typing import List, Optional
|
||||
from tenacity import retry, retry_if_exception_type, stop_after_attempt, wait_exponential
|
||||
from utils.retry import after_func
|
||||
from utils.image import image_path_to_b64
|
||||
from interfaces.image_output import ImageOutput
|
||||
|
||||
|
||||
class ImageGeneratorDoubaoSeedreamYunwuAPI:
|
||||
def __init__(
|
||||
self,
|
||||
api_key: str,
|
||||
model: str = "doubao-seedream-4-0-250828",
|
||||
|
||||
):
|
||||
self.api_key = api_key
|
||||
self.base_url = "https://yunwu.ai/v1/images/generations"
|
||||
self.model = model
|
||||
|
||||
|
||||
@retry(
|
||||
stop=stop_after_attempt(3),
|
||||
wait=wait_exponential(multiplier=1, max=30),
|
||||
retry=retry_if_exception_type((aiohttp.ClientError, asyncio.TimeoutError)),
|
||||
reraise=True,
|
||||
after=after_func,
|
||||
)
|
||||
async def generate_single_image(
|
||||
self,
|
||||
prompt: str,
|
||||
reference_image_paths: List[str] = [],
|
||||
size: Optional[str] = None,
|
||||
**kwargs,
|
||||
) -> ImageOutput:
|
||||
"""
|
||||
size: [1024x1024, 4096x4096]
|
||||
"""
|
||||
|
||||
logging.info(f"Calling {self.model} to generate image...")
|
||||
|
||||
image = [
|
||||
image_path_to_b64(path, mime=True) for path in reference_image_paths
|
||||
]
|
||||
|
||||
payload = {
|
||||
"model": self.model,
|
||||
"prompt": prompt,
|
||||
"sequential_image_generation": "disabled", # "auto" or "disabled"
|
||||
# "sequential_image_generation_options": {
|
||||
# "max_images": 1
|
||||
# },
|
||||
"response_format": "url",
|
||||
"size": size if size is not None else "1024x1024",
|
||||
}
|
||||
if len(image) > 0:
|
||||
payload["image"] = image
|
||||
|
||||
headers = {
|
||||
"Authorization": f"Bearer {self.api_key}",
|
||||
"Content-Type": "application/json",
|
||||
}
|
||||
|
||||
async with aiohttp.ClientSession() as session:
|
||||
async with session.post(self.base_url, json=payload, headers=headers) as response:
|
||||
response_json = await response.json()
|
||||
if response.status >= 400:
|
||||
raise RuntimeError(f"Image generation failed with HTTP {response.status}: {response_json}")
|
||||
|
||||
data = response_json['data'][0]['url']
|
||||
return ImageOutput(fmt="url", ext="png", data=data)
|
||||
@@ -0,0 +1,97 @@
|
||||
# https://ai.google.dev/gemini-api/docs/image-generation
|
||||
|
||||
import logging
|
||||
import asyncio
|
||||
from PIL import Image
|
||||
from typing import List, Optional
|
||||
from google import genai
|
||||
from google.genai import types
|
||||
from google.genai.errors import ClientError
|
||||
from tenacity import retry, stop_after_attempt, wait_exponential
|
||||
from interfaces.image_output import ImageOutput
|
||||
from tools.image_orientation import ensure_not_portrait, landscape_guard_requested
|
||||
from tools.image_response import image_from_response_part
|
||||
from utils.retry import after_func
|
||||
from utils.rate_limiter import RateLimiter
|
||||
|
||||
|
||||
class ImageGeneratorNanobananaGoogleAPI:
|
||||
def __init__(
|
||||
self,
|
||||
api_key: str,
|
||||
rate_limiter: Optional[RateLimiter] = None,
|
||||
):
|
||||
self.model = "gemini-2.5-flash-image"
|
||||
self.rate_limiter = rate_limiter
|
||||
self.client = genai.Client(
|
||||
api_key=api_key,
|
||||
)
|
||||
|
||||
@retry(stop=stop_after_attempt(3), wait=wait_exponential(multiplier=1, min=1, max=10), after=after_func, reraise=True)
|
||||
async def generate_single_image(
|
||||
self,
|
||||
prompt: str,
|
||||
reference_image_paths: List[str] = [],
|
||||
aspect_ratio: Optional[str] = "16:9",
|
||||
**kwargs,
|
||||
) -> ImageOutput:
|
||||
|
||||
"""
|
||||
aspect_ratio: The aspect ratio of the image.
|
||||
"""
|
||||
|
||||
logging.info(f"Calling {self.model} to generate image...")
|
||||
|
||||
# Apply rate limiting if configured
|
||||
if self.rate_limiter:
|
||||
await self.rate_limiter.acquire()
|
||||
|
||||
reference_images = [Image.open(path) for path in reference_image_paths]
|
||||
|
||||
# Retry logic for rate limit errors
|
||||
max_retries = 3
|
||||
retry_delay = 5
|
||||
|
||||
for attempt in range(max_retries):
|
||||
try:
|
||||
response = await self.client.aio.models.generate_content(
|
||||
model=self.model,
|
||||
contents=reference_images + [prompt],
|
||||
config=types.GenerateContentConfig(
|
||||
response_modalities=["IMAGE"],
|
||||
image_config=types.ImageConfig(
|
||||
aspect_ratio=aspect_ratio,
|
||||
),
|
||||
),
|
||||
)
|
||||
break
|
||||
except ClientError as e:
|
||||
if e.status_code == 429 and attempt < max_retries - 1:
|
||||
wait_time = retry_delay * (2 ** attempt)
|
||||
logging.warning(f"Rate limit hit (429), retrying in {wait_time}s... (attempt {attempt + 1}/{max_retries})")
|
||||
await asyncio.sleep(wait_time)
|
||||
else:
|
||||
raise
|
||||
|
||||
image = None
|
||||
text = ""
|
||||
for part in response.candidates[0].content.parts:
|
||||
if part.text is not None:
|
||||
text += part.text
|
||||
elif part.inline_data is not None:
|
||||
image = image_from_response_part(part)
|
||||
|
||||
if image is None:
|
||||
logging.error(f"No image generated. The response text is: {text}")
|
||||
raise ValueError("No image generated")
|
||||
|
||||
if landscape_guard_requested(
|
||||
size=kwargs.get("size"),
|
||||
aspect_ratio=aspect_ratio,
|
||||
enforce_landscape=kwargs.get("enforce_landscape", True),
|
||||
allow_portrait=kwargs.get("allow_portrait", False),
|
||||
):
|
||||
ensure_not_portrait(image)
|
||||
|
||||
return ImageOutput(fmt="pil", ext="png", data=image)
|
||||
|
||||
@@ -0,0 +1,79 @@
|
||||
# https://ai.google.dev/gemini-api/docs/image-generation?hl=zh-cn
|
||||
|
||||
import logging
|
||||
from PIL import Image
|
||||
from typing import List, Optional
|
||||
from google import genai
|
||||
from google.genai import types
|
||||
from tenacity import retry, stop_after_attempt, wait_exponential
|
||||
from interfaces.image_output import ImageOutput
|
||||
from tools.image_orientation import ensure_not_portrait, landscape_guard_requested
|
||||
from tools.image_response import image_from_response_part
|
||||
from utils.retry import after_func
|
||||
|
||||
|
||||
class ImageGeneratorNanobananaYunwuAPI:
|
||||
def __init__(
|
||||
self,
|
||||
api_key: str,
|
||||
model: str = "gemini-2.5-flash-image-preview",
|
||||
base_url: str = "https://yunwu.ai",
|
||||
):
|
||||
self.client = genai.Client(
|
||||
api_key=api_key,
|
||||
http_options=types.HttpOptions(
|
||||
base_url=base_url.rstrip("/"),
|
||||
api_version="v1beta",
|
||||
),
|
||||
)
|
||||
self.model = model
|
||||
|
||||
|
||||
@retry(stop=stop_after_attempt(3), wait=wait_exponential(multiplier=1, min=1, max=10), after=after_func, reraise=True)
|
||||
async def generate_single_image(
|
||||
self,
|
||||
prompt: str,
|
||||
reference_image_paths: List[str] = [],
|
||||
aspect_ratio: Optional[str] = "16:9",
|
||||
**kwargs,
|
||||
) -> ImageOutput:
|
||||
"""
|
||||
aspect_ratio: The aspect ratio of the image.
|
||||
"""
|
||||
|
||||
logging.info(f"Calling {self.model} to generate image...")
|
||||
|
||||
reference_images = [Image.open(path) for path in reference_image_paths]
|
||||
|
||||
response = await self.client.aio.models.generate_content(
|
||||
model=self.model,
|
||||
contents=reference_images + [prompt],
|
||||
config=types.GenerateContentConfig(
|
||||
response_modalities=["TEXT", "IMAGE"],
|
||||
image_config=types.ImageConfig(
|
||||
aspect_ratio=aspect_ratio,
|
||||
),
|
||||
),
|
||||
)
|
||||
|
||||
image = None
|
||||
text = ""
|
||||
for part in response.candidates[0].content.parts:
|
||||
if part.text is not None:
|
||||
text += part.text
|
||||
elif part.inline_data is not None:
|
||||
image = image_from_response_part(part)
|
||||
|
||||
if image is None:
|
||||
logging.error(f"No image generated. The response text is: {text}")
|
||||
raise ValueError(f"Error occurred while generating image.")
|
||||
|
||||
if landscape_guard_requested(
|
||||
size=kwargs.get("size"),
|
||||
aspect_ratio=aspect_ratio,
|
||||
enforce_landscape=kwargs.get("enforce_landscape", True),
|
||||
allow_portrait=kwargs.get("allow_portrait", False),
|
||||
):
|
||||
ensure_not_portrait(image)
|
||||
|
||||
return ImageOutput(fmt="pil", ext="png", data=image)
|
||||
@@ -0,0 +1,53 @@
|
||||
from __future__ import annotations
|
||||
|
||||
import os
|
||||
from typing import Any
|
||||
|
||||
from PIL import Image
|
||||
|
||||
|
||||
def landscape_guard_requested(*, size: Any = None, aspect_ratio: Any = None, enforce_landscape: Any = True, allow_portrait: Any = False) -> bool:
|
||||
if bool(allow_portrait):
|
||||
return False
|
||||
if bool(enforce_landscape):
|
||||
return True
|
||||
parsed = _parse_size(size)
|
||||
if parsed and parsed[0] > parsed[1]:
|
||||
return True
|
||||
parsed_ratio = _parse_size(aspect_ratio)
|
||||
return bool(parsed_ratio and parsed_ratio[0] > parsed_ratio[1])
|
||||
|
||||
|
||||
def ensure_not_portrait(image: Image.Image, *, tolerance: float | None = None) -> None:
|
||||
width, height = image.size
|
||||
if width <= 0 or height <= 0:
|
||||
return
|
||||
threshold = tolerance if tolerance is not None else _portrait_tolerance()
|
||||
if height > width * threshold:
|
||||
raise ValueError(f"Generated image is portrait-oriented ({width}x{height}); retrying for a landscape frame")
|
||||
|
||||
|
||||
def _portrait_tolerance() -> float:
|
||||
raw = os.environ.get("VIMAX_IMAGE_PORTRAIT_RETRY_TOLERANCE", "1.05")
|
||||
try:
|
||||
return max(1.0, float(raw))
|
||||
except ValueError:
|
||||
return 1.05
|
||||
|
||||
|
||||
def _parse_size(size: Any) -> tuple[int, int] | None:
|
||||
if not isinstance(size, str):
|
||||
return None
|
||||
normalized = size.lower()
|
||||
separator = "x" if "x" in normalized else ":" if ":" in normalized else ""
|
||||
if not separator:
|
||||
return None
|
||||
left, right = normalized.split(separator, 1)
|
||||
try:
|
||||
width = int(left.strip())
|
||||
height = int(right.strip())
|
||||
except ValueError:
|
||||
return None
|
||||
if width <= 0 or height <= 0:
|
||||
return None
|
||||
return width, height
|
||||
@@ -0,0 +1,40 @@
|
||||
from __future__ import annotations
|
||||
|
||||
import base64
|
||||
from io import BytesIO
|
||||
from typing import Any
|
||||
|
||||
from PIL import Image
|
||||
|
||||
|
||||
def image_from_response_part(part: Any) -> Image.Image | None:
|
||||
inline_data = getattr(part, "inline_data", None)
|
||||
if inline_data is None and isinstance(part, dict):
|
||||
inline_data = part.get("inline_data")
|
||||
if inline_data is None:
|
||||
return None
|
||||
|
||||
as_image = getattr(part, "as_image", None)
|
||||
if callable(as_image):
|
||||
image = as_image()
|
||||
if isinstance(image, Image.Image):
|
||||
return image
|
||||
|
||||
data = _value(inline_data, "data")
|
||||
if data is None:
|
||||
return None
|
||||
if isinstance(data, str):
|
||||
if data.startswith("data:") and "," in data:
|
||||
data = data.split(",", 1)[1]
|
||||
data = base64.b64decode(data)
|
||||
if isinstance(data, bytearray):
|
||||
data = bytes(data)
|
||||
if not isinstance(data, bytes):
|
||||
return None
|
||||
return Image.open(BytesIO(data)).convert("RGB")
|
||||
|
||||
|
||||
def _value(obj: Any, key: str) -> Any:
|
||||
if isinstance(obj, dict):
|
||||
return obj.get(key)
|
||||
return getattr(obj, key, None)
|
||||
@@ -0,0 +1,35 @@
|
||||
"""Structural typing contracts for rendering backends.
|
||||
|
||||
Any class that exposes the right method signatures satisfies these
|
||||
protocols -- no inheritance required. Existing generators (Google,
|
||||
Yunwu/Doubao, Yunwu/Veo) are already compliant by duck typing.
|
||||
"""
|
||||
|
||||
from typing import List, Protocol, runtime_checkable
|
||||
|
||||
from interfaces.image_output import ImageOutput
|
||||
from interfaces.video_output import VideoOutput
|
||||
|
||||
|
||||
@runtime_checkable
|
||||
class ImageGenerator(Protocol):
|
||||
"""Generates a single image from a text prompt and optional reference images."""
|
||||
|
||||
async def generate_single_image(
|
||||
self,
|
||||
prompt: str,
|
||||
reference_image_paths: List[str],
|
||||
**kwargs,
|
||||
) -> ImageOutput: ...
|
||||
|
||||
|
||||
@runtime_checkable
|
||||
class VideoGenerator(Protocol):
|
||||
"""Generates a single video from a text prompt and optional reference images."""
|
||||
|
||||
async def generate_single_video(
|
||||
self,
|
||||
prompt: str,
|
||||
reference_image_paths: List[str],
|
||||
**kwargs,
|
||||
) -> VideoOutput: ...
|
||||
@@ -0,0 +1,62 @@
|
||||
"""RenderBackend: config-driven factory for image and video generators.
|
||||
|
||||
Reads the ``image_generator`` and ``video_generator`` sections from a
|
||||
ViMax YAML config, instantiates the concrete classes via *class_path*,
|
||||
and wires up rate limiters.
|
||||
|
||||
Usage::
|
||||
|
||||
backend = RenderBackend.from_config(config)
|
||||
image = await backend.image_generator.generate_single_image(...)
|
||||
video = await backend.video_generator.generate_single_video(...)
|
||||
"""
|
||||
|
||||
import importlib
|
||||
import logging
|
||||
from dataclasses import dataclass
|
||||
from typing import Any, Dict
|
||||
|
||||
from utils.rate_limiter import RateLimiter
|
||||
|
||||
|
||||
@dataclass
|
||||
class RenderBackend:
|
||||
"""Bundles an image generator and a video generator."""
|
||||
|
||||
image_generator: Any
|
||||
video_generator: Any
|
||||
|
||||
@classmethod
|
||||
def from_config(cls, config: Dict[str, Any]) -> "RenderBackend":
|
||||
"""Build a RenderBackend from a parsed YAML config dict.
|
||||
|
||||
Rate limiters are created from ``max_requests_per_minute`` /
|
||||
``max_requests_per_day`` if present in each generator section.
|
||||
"""
|
||||
img_cfg = config["image_generator"]
|
||||
vid_cfg = config["video_generator"]
|
||||
|
||||
image_gen = _instantiate(img_cfg, _build_rate_limiter(img_cfg))
|
||||
video_gen = _instantiate(vid_cfg, _build_rate_limiter(vid_cfg))
|
||||
|
||||
logging.info("RenderBackend: image=%s, video=%s",
|
||||
img_cfg["class_path"], vid_cfg["class_path"])
|
||||
|
||||
return cls(image_generator=image_gen, video_generator=video_gen)
|
||||
|
||||
|
||||
def _build_rate_limiter(section: Dict[str, Any]) -> RateLimiter | None:
|
||||
rpm = section.get("max_requests_per_minute")
|
||||
rpd = section.get("max_requests_per_day")
|
||||
if rpm or rpd:
|
||||
return RateLimiter(max_requests_per_minute=rpm, max_requests_per_day=rpd)
|
||||
return None
|
||||
|
||||
|
||||
def _instantiate(section: Dict[str, Any], rate_limiter: RateLimiter | None) -> Any:
|
||||
module_path, cls_name = section["class_path"].rsplit(".", 1)
|
||||
cls = getattr(importlib.import_module(module_path), cls_name)
|
||||
init_args = dict(section.get("init_args", {}))
|
||||
if rate_limiter is not None:
|
||||
init_args["rate_limiter"] = rate_limiter
|
||||
return cls(**init_args)
|
||||
@@ -0,0 +1,83 @@
|
||||
from typing import List
|
||||
import aiohttp
|
||||
import asyncio
|
||||
from tenacity import retry, retry_if_exception_type, stop_after_attempt, wait_exponential
|
||||
import logging
|
||||
|
||||
|
||||
class RerankerBgeSiliconapi:
|
||||
def __init__(
|
||||
self,
|
||||
api_key: str,
|
||||
base_url: str,
|
||||
model: str = "BAAI/bge-reranker-v2-m3",
|
||||
):
|
||||
self.api_key = api_key
|
||||
self.base_url = base_url
|
||||
self.model = model
|
||||
# return_documents: bool = True,
|
||||
|
||||
|
||||
@retry(
|
||||
stop=stop_after_attempt(3),
|
||||
wait=wait_exponential(multiplier=1, max=30),
|
||||
retry=retry_if_exception_type((aiohttp.ClientError, asyncio.TimeoutError)),
|
||||
reraise=True,
|
||||
after=lambda retry_state: logging.warning(f"Retrying SiliconReranker due to error: {retry_state.outcome.exception()}"),
|
||||
)
|
||||
async def __call__(
|
||||
self,
|
||||
documents: List[str],
|
||||
query: str,
|
||||
top_n: int,
|
||||
) -> List[str]:
|
||||
|
||||
url = f"{self.base_url}/rerank"
|
||||
|
||||
payload = {
|
||||
"model": self.model,
|
||||
"query": query,
|
||||
"documents": documents,
|
||||
"top_n": top_n,
|
||||
"return_documents": True,
|
||||
}
|
||||
|
||||
|
||||
headers = {
|
||||
'Accept': 'application/json',
|
||||
'Authorization': f'Bearer {self.api_key}',
|
||||
'Content-Type': 'application/json'
|
||||
}
|
||||
|
||||
async with aiohttp.ClientSession() as session:
|
||||
async with session.post(url, json=payload, headers=headers) as resp:
|
||||
response = await resp.json()
|
||||
if resp.status >= 400:
|
||||
raise RuntimeError(f"Rerank request failed with HTTP {resp.status}: {response}")
|
||||
|
||||
|
||||
"""
|
||||
{
|
||||
"id": "<string>",
|
||||
"results": [
|
||||
{
|
||||
"document": {
|
||||
"text": "<string>"
|
||||
},
|
||||
"index": 123,
|
||||
"relevance_score": 123
|
||||
}
|
||||
],
|
||||
"tokens": {
|
||||
"input_tokens": 123,
|
||||
"output_tokens": 123
|
||||
}
|
||||
}
|
||||
"""
|
||||
|
||||
results = []
|
||||
|
||||
for result in response["results"]:
|
||||
results.append((result["document"]["text"], result["relevance_score"]))
|
||||
|
||||
return results
|
||||
@@ -0,0 +1,212 @@
|
||||
import logging
|
||||
from typing import List, Literal
|
||||
import asyncio
|
||||
import aiohttp
|
||||
from interfaces.video_output import VideoOutput
|
||||
from utils.image import image_path_to_b64
|
||||
|
||||
|
||||
class VideoGeneratorDoubaoSeedanceYunwuAPI:
|
||||
def __init__(
|
||||
self,
|
||||
api_key: str,
|
||||
t2v_model: str = "doubao-seedance-1-0-lite-t2v-250428",
|
||||
ff2v_model: str = "doubao-seedance-1-0-lite-i2v-250428",
|
||||
flf2v_model: str = "doubao-seedance-1-0-lite-i2v-250428",
|
||||
max_create_attempts: int = 3,
|
||||
poll_interval: int = 2,
|
||||
max_poll_attempts: int = 300,
|
||||
):
|
||||
self.api_key = api_key
|
||||
self.t2v_model = t2v_model
|
||||
self.ff2v_model = ff2v_model
|
||||
self.flf2v_model = flf2v_model
|
||||
self.max_create_attempts = max_create_attempts
|
||||
self.poll_interval = poll_interval
|
||||
self.max_poll_attempts = max_poll_attempts
|
||||
|
||||
|
||||
async def create_video_generation_task(
|
||||
self,
|
||||
prompt: str,
|
||||
reference_image_paths: List[str],
|
||||
resolution: Literal["480p", "720p", "1080p"] = "720p",
|
||||
aspect_ratio: str = "16:9",
|
||||
fps: Literal[16, 24] = 16,
|
||||
duration: Literal[5, 10] = 5,
|
||||
) -> str:
|
||||
"""
|
||||
Create a video generation task and return the task ID.
|
||||
|
||||
Args:
|
||||
prompt: Text prompt for video generation
|
||||
reference_image_paths: List of 1 or 2 reference images
|
||||
|
||||
Returns:
|
||||
Task ID string
|
||||
"""
|
||||
if len(reference_image_paths) == 0:
|
||||
model = self.t2v_model
|
||||
elif len(reference_image_paths) == 1:
|
||||
model = self.ff2v_model
|
||||
elif len(reference_image_paths) == 2:
|
||||
model = self.flf2v_model
|
||||
else:
|
||||
raise ValueError("reference_image_paths must contain 1 or 2 images.")
|
||||
|
||||
logging.info(f"Calling {model} to generate video...")
|
||||
|
||||
url = "https://yunwu.ai/volc/v1/contents/generations/tasks"
|
||||
|
||||
|
||||
content = [
|
||||
{
|
||||
"type": "text",
|
||||
"text": prompt + f" --rs {resolution} --rt {aspect_ratio} --dur {duration} --fps {fps} --wm false --seed -1 --cf false"
|
||||
}
|
||||
]
|
||||
if len(reference_image_paths) >= 1:
|
||||
content.append(
|
||||
{
|
||||
"type": "image_url",
|
||||
"image_url": {
|
||||
"url": image_path_to_b64(reference_image_paths[0])
|
||||
},
|
||||
"role": "first_frame",
|
||||
}
|
||||
)
|
||||
if len(reference_image_paths) >= 2:
|
||||
content.append(
|
||||
{
|
||||
"type": "image_url",
|
||||
"image_url": {
|
||||
"url": image_path_to_b64(reference_image_paths[1])
|
||||
},
|
||||
"role": "last_frame",
|
||||
}
|
||||
)
|
||||
|
||||
payload = {
|
||||
"model": model,
|
||||
"content": content
|
||||
}
|
||||
|
||||
headers = {
|
||||
'Authorization': f'Bearer {self.api_key}',
|
||||
'Content-Type': 'application/json'
|
||||
}
|
||||
|
||||
last_error = None
|
||||
for attempt in range(1, self.max_create_attempts + 1):
|
||||
try:
|
||||
async with aiohttp.ClientSession() as session:
|
||||
async with session.post(url, headers=headers, json=payload) as response:
|
||||
response_json = await response.json()
|
||||
http_status = response.status
|
||||
logging.debug(f"Response: {response_json}")
|
||||
except Exception as e:
|
||||
last_error = e
|
||||
logging.error(f"Error occurred while creating video generation task (attempt {attempt}/{self.max_create_attempts}): {e}")
|
||||
if attempt < self.max_create_attempts:
|
||||
await asyncio.sleep(attempt)
|
||||
continue
|
||||
|
||||
if http_status >= 400:
|
||||
message = f"Video generation task creation failed with HTTP {http_status}: {response_json}"
|
||||
if http_status < 500:
|
||||
raise RuntimeError(message)
|
||||
last_error = RuntimeError(message)
|
||||
logging.error(f"{message} (attempt {attempt}/{self.max_create_attempts})")
|
||||
if attempt < self.max_create_attempts:
|
||||
await asyncio.sleep(attempt)
|
||||
continue
|
||||
|
||||
task_id = response_json.get("id")
|
||||
if not task_id:
|
||||
raise RuntimeError(f"Video generation task creation returned no task id: {response_json}")
|
||||
logging.info(f"Video generation task created successfully. Task ID: {task_id}")
|
||||
return task_id
|
||||
|
||||
raise RuntimeError(f"Failed to create video generation task after {self.max_create_attempts} attempts.") from last_error
|
||||
|
||||
async def query_video_generation_task(
|
||||
self,
|
||||
task_id: str,
|
||||
) -> str:
|
||||
"""
|
||||
Query the video generation task until completion and return the video URL.
|
||||
|
||||
Args:
|
||||
task_id: Task ID to query
|
||||
|
||||
Returns:
|
||||
Video URL string
|
||||
"""
|
||||
url = f"https://yunwu.ai/volc/v1/contents/generations/tasks/{task_id}"
|
||||
headers = {
|
||||
'Authorization': f'Bearer {self.api_key}',
|
||||
}
|
||||
|
||||
attempts = 0
|
||||
consecutive_errors = 0
|
||||
while True:
|
||||
if attempts >= self.max_poll_attempts:
|
||||
raise TimeoutError(f"Video generation did not complete after {attempts} polls.")
|
||||
attempts += 1
|
||||
|
||||
try:
|
||||
async with aiohttp.ClientSession() as session:
|
||||
async with session.get(url, headers=headers) as response:
|
||||
response_json = await response.json()
|
||||
http_status = response.status
|
||||
except Exception as e:
|
||||
consecutive_errors += 1
|
||||
if consecutive_errors >= 5:
|
||||
raise RuntimeError(f"Querying video generation task failed {consecutive_errors} times in a row.") from e
|
||||
logging.error(f"Error occurred while querying video generation task: {e}. Retrying in {self.poll_interval} seconds...")
|
||||
await asyncio.sleep(self.poll_interval)
|
||||
continue
|
||||
consecutive_errors = 0
|
||||
|
||||
if http_status >= 400:
|
||||
raise RuntimeError(f"Querying video generation task failed with HTTP {http_status}: {response_json}")
|
||||
|
||||
status = response_json.get("status")
|
||||
if status == "succeeded":
|
||||
video_url = response_json["content"]["video_url"]
|
||||
logging.info(f"Video generation completed successfully. Video URL: {video_url}")
|
||||
return video_url
|
||||
elif status == "failed":
|
||||
logging.error(f"Video generation failed. Response: {response_json}")
|
||||
raise ValueError("Video generation failed.")
|
||||
else:
|
||||
logging.info(f"Video generation is still in progress. Checking again in {self.poll_interval} seconds...")
|
||||
await asyncio.sleep(self.poll_interval)
|
||||
|
||||
async def generate_single_video(
|
||||
self,
|
||||
prompt: str,
|
||||
reference_image_paths: List[str],
|
||||
resolution: Literal["480p", "720p", "1080p"] = "720p",
|
||||
aspect_ratio: str = "16:9",
|
||||
fps: Literal[16, 24] = 16,
|
||||
duration: Literal[5, 10] = 5,
|
||||
**kwargs,
|
||||
) -> VideoOutput:
|
||||
"""
|
||||
Generate a single video by creating a task and waiting for completion.
|
||||
|
||||
Args:
|
||||
prompt: Text prompt for video generation
|
||||
reference_image_paths: List of 1 or 2 reference images
|
||||
resolution: Resolution of the video
|
||||
aspect_ratio: Aspect ratio of the video
|
||||
fps: Frames per second of the video
|
||||
duration: Duration of the video
|
||||
Returns:
|
||||
VideoOutput containing the video URL
|
||||
"""
|
||||
task_id = await self.create_video_generation_task(prompt, reference_image_paths, resolution, aspect_ratio, fps, duration)
|
||||
video_url = await self.query_video_generation_task(task_id)
|
||||
return VideoOutput(fmt="url", ext="mp4", data=video_url)
|
||||
|
||||
@@ -0,0 +1,224 @@
|
||||
import asyncio
|
||||
import logging
|
||||
from typing import List, Optional
|
||||
|
||||
import aiohttp
|
||||
|
||||
from interfaces.video_output import VideoOutput
|
||||
from utils.image import image_path_to_b64
|
||||
from utils.rate_limiter import RateLimiter
|
||||
|
||||
|
||||
class VideoGeneratorOmniYunwuAPI:
|
||||
def __init__(
|
||||
self,
|
||||
api_key: str,
|
||||
t2v_model: str = "omni-flash",
|
||||
i2v_model: str = "omni-flash",
|
||||
base_url: str = "https://yunwu.ai",
|
||||
seconds: int = 8,
|
||||
enable_upsample: bool = False,
|
||||
enable_sample: Optional[bool] = None,
|
||||
poll_interval: int = 2,
|
||||
max_poll_attempts: Optional[int] = 300,
|
||||
max_create_attempts: int = 3,
|
||||
rate_limiter: Optional[RateLimiter] = None,
|
||||
):
|
||||
self.api_key = api_key
|
||||
self.t2v_model = t2v_model
|
||||
self.i2v_model = i2v_model
|
||||
self.base_url = base_url.rstrip("/")
|
||||
self.seconds = seconds
|
||||
self.enable_upsample = enable_upsample
|
||||
self.enable_sample = enable_sample
|
||||
self.poll_interval = poll_interval
|
||||
self.max_poll_attempts = max_poll_attempts
|
||||
self.max_create_attempts = max_create_attempts
|
||||
self.rate_limiter = rate_limiter
|
||||
|
||||
def _headers(self) -> dict:
|
||||
return {
|
||||
"Accept": "application/json",
|
||||
"Authorization": f"Bearer {self.api_key}",
|
||||
"Content-Type": "application/json",
|
||||
}
|
||||
|
||||
def _image_uri(self, image_path: str) -> str:
|
||||
if image_path.startswith(("http://", "https://", "data:")):
|
||||
return image_path
|
||||
return image_path_to_b64(image_path, mime=True)
|
||||
|
||||
def _build_payload(
|
||||
self,
|
||||
prompt: str,
|
||||
reference_image_paths: List[str],
|
||||
aspect_ratio: str,
|
||||
seconds: Optional[int],
|
||||
size: Optional[str],
|
||||
enable_upsample: Optional[bool],
|
||||
enable_sample: Optional[bool],
|
||||
) -> dict:
|
||||
if len(reference_image_paths) > 3:
|
||||
raise ValueError("The number of reference images must be no more than 3")
|
||||
|
||||
payload = {
|
||||
"model": self.t2v_model if len(reference_image_paths) == 0 else self.i2v_model,
|
||||
"prompt": prompt,
|
||||
"seconds": str(seconds or self.seconds),
|
||||
}
|
||||
|
||||
if len(reference_image_paths) == 0:
|
||||
payload["type"] = 1
|
||||
elif len(reference_image_paths) <= 2:
|
||||
payload["type"] = 2
|
||||
payload["images"] = [self._image_uri(path) for path in reference_image_paths]
|
||||
else:
|
||||
payload["type"] = 3
|
||||
payload["images"] = [self._image_uri(path) for path in reference_image_paths]
|
||||
|
||||
if aspect_ratio:
|
||||
payload["aspect_ratio"] = aspect_ratio
|
||||
if size:
|
||||
payload["size"] = size
|
||||
if enable_upsample is not None:
|
||||
payload["enable_upsample"] = enable_upsample
|
||||
if enable_sample is not None:
|
||||
payload["enable_sample"] = enable_sample
|
||||
|
||||
return payload
|
||||
|
||||
async def create_video_generation_task(
|
||||
self,
|
||||
prompt: str,
|
||||
reference_image_paths: List[str],
|
||||
aspect_ratio: str = "16:9",
|
||||
seconds: Optional[int] = None,
|
||||
size: Optional[str] = None,
|
||||
enable_upsample: Optional[bool] = None,
|
||||
enable_sample: Optional[bool] = None,
|
||||
) -> tuple[str, str]:
|
||||
payload = self._build_payload(
|
||||
prompt=prompt,
|
||||
reference_image_paths=reference_image_paths,
|
||||
aspect_ratio=aspect_ratio,
|
||||
seconds=seconds,
|
||||
size=size,
|
||||
enable_upsample=self.enable_upsample if enable_upsample is None else enable_upsample,
|
||||
enable_sample=self.enable_sample if enable_sample is None else enable_sample,
|
||||
)
|
||||
|
||||
logging.info("Calling %s to generate video...", payload["model"])
|
||||
|
||||
if self.rate_limiter:
|
||||
await self.rate_limiter.acquire()
|
||||
|
||||
url = f"{self.base_url}/v1/video/create"
|
||||
last_error = None
|
||||
for attempt in range(1, self.max_create_attempts + 1):
|
||||
try:
|
||||
async with aiohttp.ClientSession() as session:
|
||||
async with session.post(url, headers=self._headers(), json=payload) as response:
|
||||
response_json = await response.json()
|
||||
http_status = response.status
|
||||
logging.debug("Response: %s", response_json)
|
||||
except Exception as e:
|
||||
last_error = e
|
||||
logging.error(
|
||||
"Error occurred while creating video generation task (attempt %s/%s): %s",
|
||||
attempt,
|
||||
self.max_create_attempts,
|
||||
e,
|
||||
)
|
||||
if attempt < self.max_create_attempts:
|
||||
await asyncio.sleep(attempt)
|
||||
continue
|
||||
|
||||
if http_status >= 400:
|
||||
message = f"Video generation task creation failed with HTTP {http_status}: {response_json}"
|
||||
if http_status < 500:
|
||||
raise RuntimeError(message)
|
||||
last_error = RuntimeError(message)
|
||||
logging.error("%s (attempt %s/%s)", message, attempt, self.max_create_attempts)
|
||||
if attempt < self.max_create_attempts:
|
||||
await asyncio.sleep(attempt)
|
||||
continue
|
||||
|
||||
task_id = response_json.get("id")
|
||||
if not task_id:
|
||||
raise RuntimeError(f"Video generation task creation returned no task id: {response_json}")
|
||||
logging.info("Video generation task created successfully. Task ID: %s", task_id)
|
||||
return task_id, payload["model"]
|
||||
|
||||
raise RuntimeError(
|
||||
f"Failed to create video generation task after {self.max_create_attempts} attempts."
|
||||
) from last_error
|
||||
|
||||
async def query_video_generation_task(self, task_id: str, model: str) -> str:
|
||||
url = f"{self.base_url}/v1/video/query"
|
||||
params = {"id": task_id, "model": model}
|
||||
|
||||
attempts = 0
|
||||
while True:
|
||||
if self.max_poll_attempts is not None and attempts >= self.max_poll_attempts:
|
||||
raise TimeoutError(f"Video generation did not complete after {attempts} polls.")
|
||||
attempts += 1
|
||||
|
||||
try:
|
||||
async with aiohttp.ClientSession() as session:
|
||||
async with session.get(url, headers=self._headers(), params=params) as response:
|
||||
response_json = await response.json()
|
||||
logging.debug("Response: %s", response_json)
|
||||
except Exception as e:
|
||||
logging.error(
|
||||
"Error occurred while querying video generation task: %s. Retrying in %s seconds...",
|
||||
e,
|
||||
self.poll_interval,
|
||||
)
|
||||
await asyncio.sleep(self.poll_interval)
|
||||
continue
|
||||
|
||||
status = response_json.get("status")
|
||||
if status == "completed":
|
||||
detail = response_json.get("detail") or {}
|
||||
video_url = (
|
||||
response_json.get("video_url")
|
||||
or detail.get("upsample_video_url")
|
||||
or detail.get("video_url")
|
||||
)
|
||||
if not video_url:
|
||||
raise RuntimeError(f"Video generation completed without a video URL: {response_json}")
|
||||
logging.info("Video generation completed successfully. Video URL: %s", video_url)
|
||||
return video_url
|
||||
|
||||
if status in {"failed", "error"}:
|
||||
raise RuntimeError(f"Video generation failed: {response_json}")
|
||||
|
||||
logging.info("Video generation status: %s, waiting %s seconds...", status, self.poll_interval)
|
||||
await asyncio.sleep(self.poll_interval)
|
||||
|
||||
async def generate_single_video(
|
||||
self,
|
||||
prompt: str,
|
||||
reference_image_paths: List[str],
|
||||
aspect_ratio: str = "16:9",
|
||||
seconds: Optional[int] = None,
|
||||
size: Optional[str] = None,
|
||||
enable_upsample: Optional[bool] = None,
|
||||
enable_sample: Optional[bool] = None,
|
||||
**kwargs,
|
||||
) -> VideoOutput:
|
||||
task_id, model = await self.create_video_generation_task(
|
||||
prompt=prompt,
|
||||
reference_image_paths=reference_image_paths,
|
||||
aspect_ratio=aspect_ratio,
|
||||
seconds=seconds,
|
||||
size=size,
|
||||
enable_upsample=enable_upsample,
|
||||
enable_sample=enable_sample,
|
||||
)
|
||||
video_url = await self.query_video_generation_task(task_id, model)
|
||||
return VideoOutput(fmt="url", ext="mp4", data=video_url)
|
||||
|
||||
|
||||
class VideoGeneratorOminiYunwuAPI(VideoGeneratorOmniYunwuAPI):
|
||||
"""Backward-compatible alias for the common "omini" spelling."""
|
||||
@@ -0,0 +1,201 @@
|
||||
import asyncio
|
||||
import logging
|
||||
import os
|
||||
from typing import List
|
||||
from urllib.parse import urljoin
|
||||
|
||||
import aiohttp
|
||||
|
||||
from interfaces.video_output import VideoOutput
|
||||
from utils.image import image_path_to_b64
|
||||
|
||||
|
||||
def _env_int(name: str, default: int) -> int:
|
||||
try:
|
||||
return max(0, int(os.environ.get(name, str(default))))
|
||||
except ValueError:
|
||||
return default
|
||||
|
||||
|
||||
def _env_float(name: str, default: float) -> float:
|
||||
try:
|
||||
return max(0.0, float(os.environ.get(name, str(default))))
|
||||
except ValueError:
|
||||
return default
|
||||
|
||||
|
||||
def _env_bool(name: str, default: bool) -> bool:
|
||||
raw = os.environ.get(name)
|
||||
if raw is None:
|
||||
return default
|
||||
return raw.strip().lower() in {"1", "true", "yes", "on"}
|
||||
|
||||
|
||||
def _emit_progress(progress, stage: str, message: str, metadata: dict | None = None) -> None:
|
||||
if progress is not None:
|
||||
progress(stage, message, metadata or {})
|
||||
|
||||
|
||||
class VideoGeneratorOpenRouterAPI:
|
||||
def __init__(
|
||||
self,
|
||||
api_key: str,
|
||||
model: str = "google/veo-3.1-lite",
|
||||
base_url: str = "https://openrouter.ai/api/v1",
|
||||
http_referer: str = "",
|
||||
app_title: str = "ViMax",
|
||||
):
|
||||
self.api_key = api_key
|
||||
self.model = model
|
||||
self.base_url = base_url.rstrip("/")
|
||||
self.http_referer = http_referer
|
||||
self.app_title = app_title
|
||||
|
||||
async def generate_single_video(
|
||||
self,
|
||||
prompt: str = "",
|
||||
reference_image_paths: List[str] = [],
|
||||
aspect_ratio: str = "16:9",
|
||||
**kwargs,
|
||||
) -> VideoOutput:
|
||||
progress = kwargs.get("progress")
|
||||
request_timeout_seconds = _env_float("VIMAX_VIDEO_REQUEST_TIMEOUT_SECONDS", 60.0)
|
||||
query_timeout_seconds = _env_float("VIMAX_VIDEO_QUERY_TIMEOUT_SECONDS", 600.0)
|
||||
poll_interval_seconds = _env_float("VIMAX_VIDEO_POLL_INTERVAL_SECONDS", 10.0)
|
||||
duration = _env_int("VIMAX_OPENROUTER_VIDEO_DURATION", 8)
|
||||
resolution = os.environ.get("VIMAX_OPENROUTER_VIDEO_RESOLUTION", "720p")
|
||||
generate_audio = _env_bool("VIMAX_OPENROUTER_GENERATE_AUDIO", True)
|
||||
|
||||
payload = {
|
||||
"model": self.model,
|
||||
"prompt": prompt,
|
||||
"aspect_ratio": aspect_ratio,
|
||||
"duration": duration,
|
||||
"resolution": resolution,
|
||||
"generate_audio": generate_audio,
|
||||
}
|
||||
frame_images = _frame_images(reference_image_paths)
|
||||
if frame_images:
|
||||
payload["frame_images"] = frame_images
|
||||
|
||||
headers = self._headers()
|
||||
timeout = aiohttp.ClientTimeout(total=request_timeout_seconds)
|
||||
_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)})
|
||||
|
||||
create_status, create_payload = await _post_json(
|
||||
f"{self.base_url}/videos",
|
||||
headers=headers,
|
||||
payload=payload,
|
||||
timeout=timeout,
|
||||
hard_timeout_seconds=request_timeout_seconds,
|
||||
)
|
||||
if create_status >= 400:
|
||||
raise RuntimeError(f"OpenRouter video create failed with HTTP {create_status}: {create_payload}")
|
||||
job_id = create_payload.get("id")
|
||||
polling_url = create_payload.get("polling_url")
|
||||
if not job_id or not polling_url:
|
||||
raise RuntimeError(f"OpenRouter video create response missing id or polling_url: {create_payload}")
|
||||
_emit_progress(progress, "video_task_created", "OpenRouter video generation task created", {"model": self.model, "job_id": job_id, "status": create_payload.get("status")})
|
||||
|
||||
poll_url = _absolute_url(self.base_url, polling_url)
|
||||
deadline = asyncio.get_running_loop().time() + query_timeout_seconds if query_timeout_seconds > 0 else None
|
||||
last_status = create_payload.get("status")
|
||||
last_payload = create_payload
|
||||
while deadline is None or asyncio.get_running_loop().time() < deadline:
|
||||
await asyncio.sleep(poll_interval_seconds)
|
||||
poll_status, poll_payload = await _get_json(
|
||||
poll_url,
|
||||
headers=headers,
|
||||
timeout=timeout,
|
||||
hard_timeout_seconds=request_timeout_seconds,
|
||||
)
|
||||
if poll_status >= 400:
|
||||
raise RuntimeError(f"OpenRouter video poll failed with HTTP {poll_status}: {poll_payload}")
|
||||
last_payload = poll_payload
|
||||
status = poll_payload.get("status")
|
||||
last_status = status
|
||||
_emit_progress(progress, "video_status", f"OpenRouter video generation status: {status}", {"model": self.model, "job_id": job_id, "status": status})
|
||||
|
||||
if status == "completed":
|
||||
urls = poll_payload.get("unsigned_urls") or []
|
||||
if urls:
|
||||
content_url = urls[0]
|
||||
else:
|
||||
content_url = f"{self.base_url}/videos/{job_id}/content?index=0"
|
||||
_emit_progress(progress, "video_download_start", "Downloading OpenRouter video output", {"model": self.model, "job_id": job_id})
|
||||
download_status, data = await _get_bytes(
|
||||
content_url,
|
||||
headers=headers if _needs_authorization(content_url) else {},
|
||||
timeout=timeout,
|
||||
hard_timeout_seconds=request_timeout_seconds,
|
||||
)
|
||||
if download_status >= 400:
|
||||
raise RuntimeError(f"OpenRouter video content download failed with HTTP {download_status}: {data[:500]!r}")
|
||||
_emit_progress(progress, "video_completed", "OpenRouter video generation completed and downloaded", {"model": self.model, "job_id": job_id})
|
||||
return VideoOutput(fmt="bytes", ext="mp4", data=data)
|
||||
if status in {"failed", "cancelled", "expired"}:
|
||||
raise RuntimeError(f"OpenRouter video generation {status} for job {job_id}: {poll_payload.get('error') or poll_payload}")
|
||||
|
||||
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}")
|
||||
|
||||
def _headers(self) -> dict[str, str]:
|
||||
headers = {
|
||||
"Authorization": f"Bearer {self.api_key}",
|
||||
"Content-Type": "application/json",
|
||||
}
|
||||
if self.http_referer:
|
||||
headers["HTTP-Referer"] = self.http_referer
|
||||
if self.app_title:
|
||||
headers["X-OpenRouter-Title"] = self.app_title
|
||||
return headers
|
||||
|
||||
|
||||
def _frame_images(reference_image_paths: List[str]) -> list[dict]:
|
||||
if len(reference_image_paths) > 2:
|
||||
raise ValueError("OpenRouter video generation supports at most first and last frame images")
|
||||
frame_types = ["first_frame", "last_frame"]
|
||||
return [
|
||||
{
|
||||
"type": "image_url",
|
||||
"image_url": {"url": image_path_to_b64(path, mime=True)},
|
||||
"frame_type": frame_types[index],
|
||||
}
|
||||
for index, path in enumerate(reference_image_paths)
|
||||
]
|
||||
|
||||
|
||||
def _absolute_url(base_url: str, url: str) -> str:
|
||||
if url.startswith("http://") or url.startswith("https://"):
|
||||
return url
|
||||
return urljoin(f"{base_url.rstrip('/')}/", url.lstrip("/"))
|
||||
|
||||
|
||||
def _needs_authorization(url: str) -> bool:
|
||||
return url.startswith("https://openrouter.ai/api/")
|
||||
|
||||
|
||||
async def _post_json(url: str, *, headers: dict[str, str], payload: dict, timeout: aiohttp.ClientTimeout, hard_timeout_seconds: float) -> tuple[int, dict]:
|
||||
async def request() -> tuple[int, dict]:
|
||||
async with aiohttp.ClientSession(timeout=timeout) as session:
|
||||
async with session.post(url, headers=headers, json=payload) as response:
|
||||
return response.status, await response.json(content_type=None)
|
||||
|
||||
return await asyncio.wait_for(request(), timeout=hard_timeout_seconds + 5)
|
||||
|
||||
|
||||
async def _get_json(url: str, *, headers: dict[str, str], timeout: aiohttp.ClientTimeout, hard_timeout_seconds: float) -> tuple[int, dict]:
|
||||
async def request() -> tuple[int, dict]:
|
||||
async with aiohttp.ClientSession(timeout=timeout) as session:
|
||||
async with session.get(url, headers=headers) as response:
|
||||
return response.status, await response.json(content_type=None)
|
||||
|
||||
return await asyncio.wait_for(request(), timeout=hard_timeout_seconds + 5)
|
||||
|
||||
|
||||
async def _get_bytes(url: str, *, headers: dict[str, str], timeout: aiohttp.ClientTimeout, hard_timeout_seconds: float) -> tuple[int, bytes]:
|
||||
async def request() -> tuple[int, bytes]:
|
||||
async with aiohttp.ClientSession(timeout=timeout) as session:
|
||||
async with session.get(url, headers=headers) as response:
|
||||
return response.status, await response.read()
|
||||
|
||||
return await asyncio.wait_for(request(), timeout=hard_timeout_seconds + 5)
|
||||
@@ -0,0 +1,116 @@
|
||||
import logging
|
||||
from typing import List, Optional
|
||||
import asyncio
|
||||
from google import genai
|
||||
from google.genai import types
|
||||
from google.genai.errors import ClientError
|
||||
from interfaces.video_output import VideoOutput
|
||||
from utils.rate_limiter import RateLimiter
|
||||
|
||||
# https://ai.google.dev/gemini-api/docs/video-generation?hl=zh-cn
|
||||
|
||||
|
||||
class VideoGeneratorVeoGoogleAPI:
|
||||
def __init__(
|
||||
self,
|
||||
api_key: str,
|
||||
t2v_model: str = "veo-3.1-generate-preview",
|
||||
ff2v_model: str = "veo-3.1-generate-preview",
|
||||
flf2v_model: str = "veo-3.1-generate-preview",
|
||||
rate_limiter: Optional[RateLimiter] = None,
|
||||
):
|
||||
self.api_key = api_key
|
||||
self.t2v_model = t2v_model
|
||||
self.ff2v_model = ff2v_model
|
||||
self.flf2v_model = flf2v_model
|
||||
self.rate_limiter = rate_limiter
|
||||
|
||||
self.client = genai.Client(
|
||||
api_key=api_key,
|
||||
)
|
||||
|
||||
async def generate_single_video(
|
||||
self,
|
||||
prompt: str,
|
||||
reference_image_paths: List[str],
|
||||
resolution: str = "1080p",
|
||||
aspect_ratio: str = "16:9",
|
||||
duration: int = 8,
|
||||
**kwargs,
|
||||
) -> VideoOutput:
|
||||
|
||||
params = {
|
||||
"prompt": prompt,
|
||||
}
|
||||
config_params = {
|
||||
"resolution": resolution,
|
||||
"aspect_ratio": aspect_ratio,
|
||||
"duration_seconds": duration,
|
||||
}
|
||||
if len(reference_image_paths) == 0:
|
||||
params["model"] = self.t2v_model
|
||||
elif len(reference_image_paths) == 1:
|
||||
params["model"] = self.ff2v_model
|
||||
params["image"] = types.Image.from_file(location=reference_image_paths[0])
|
||||
elif len(reference_image_paths) == 2:
|
||||
params["model"] = self.flf2v_model
|
||||
params["image"] = types.Image.from_file(location=reference_image_paths[0])
|
||||
config_params["last_frame"] = types.Image.from_file(location=reference_image_paths[1])
|
||||
else:
|
||||
raise ValueError("The number of reference images must be no more than 2")
|
||||
|
||||
logging.info(f"Calling {params['model']} to generate video...")
|
||||
|
||||
# Apply rate limiting if configured
|
||||
if self.rate_limiter:
|
||||
await self.rate_limiter.acquire()
|
||||
|
||||
# Retry logic for rate limit errors
|
||||
max_retries = 3
|
||||
retry_delay = 5
|
||||
|
||||
for attempt in range(max_retries):
|
||||
try:
|
||||
operation = self.client.models.generate_videos(
|
||||
**params,
|
||||
config=types.GenerateVideosConfig(**config_params),
|
||||
)
|
||||
break
|
||||
except ClientError as e:
|
||||
if e.status_code == 429 and attempt < max_retries - 1:
|
||||
wait_time = retry_delay * (2 ** attempt)
|
||||
logging.warning(f"Rate limit hit (429), retrying in {wait_time}s... (attempt {attempt + 1}/{max_retries})")
|
||||
await asyncio.sleep(wait_time)
|
||||
else:
|
||||
raise
|
||||
|
||||
while not operation.done:
|
||||
await asyncio.sleep(2)
|
||||
operation = self.client.operations.get(operation)
|
||||
logging.info(f"Video generation not completed, waiting 2 seconds...")
|
||||
|
||||
# Check if operation completed successfully
|
||||
if operation.error:
|
||||
error_msg = f"Video generation failed: {operation.error}"
|
||||
logging.error(error_msg)
|
||||
raise RuntimeError(error_msg)
|
||||
|
||||
if not operation.response:
|
||||
error_msg = "Video generation completed but no response received"
|
||||
logging.error(error_msg)
|
||||
raise RuntimeError(error_msg)
|
||||
|
||||
if not hasattr(operation.response, 'generated_videos') or not operation.response.generated_videos:
|
||||
error_msg = "Video generation completed but no videos were generated"
|
||||
logging.error(error_msg)
|
||||
raise RuntimeError(error_msg)
|
||||
|
||||
generated_video = operation.response.generated_videos[0]
|
||||
self.client.files.download(file=generated_video.video)
|
||||
|
||||
video_output = VideoOutput(
|
||||
fmt="bytes",
|
||||
ext="mp4",
|
||||
data=generated_video.video.video_bytes,
|
||||
)
|
||||
return video_output
|
||||
@@ -0,0 +1,177 @@
|
||||
import logging
|
||||
from typing import List, Optional
|
||||
from PIL import Image
|
||||
import asyncio
|
||||
import aiohttp
|
||||
import os
|
||||
from interfaces.video_output import VideoOutput
|
||||
from utils.image import image_path_to_b64
|
||||
|
||||
|
||||
def _env_int(name: str, default: int) -> int:
|
||||
try:
|
||||
return max(0, int(os.environ.get(name, str(default))))
|
||||
except ValueError:
|
||||
return default
|
||||
|
||||
|
||||
def _env_float(name: str, default: float) -> float:
|
||||
try:
|
||||
return max(0.0, float(os.environ.get(name, str(default))))
|
||||
except ValueError:
|
||||
return default
|
||||
|
||||
|
||||
def _emit_progress(progress, stage: str, message: str, metadata: dict | None = None) -> None:
|
||||
if progress is not None:
|
||||
progress(stage, message, metadata or {})
|
||||
|
||||
|
||||
class VideoGeneratorVeoYunwuAPI:
|
||||
def __init__(
|
||||
self,
|
||||
api_key: str,
|
||||
t2v_model: str = "veo3.1-fast", # text to video
|
||||
ff2v_model: str = "veo3.1-fast", # first frame to video
|
||||
flf2v_model: str = "veo2-fast-frames", # first and last frame to video
|
||||
base_url: str = "https://yunwu.ai",
|
||||
):
|
||||
"""
|
||||
all models:
|
||||
veo2
|
||||
veo2-fast
|
||||
veo2-fast-frames
|
||||
veo2-fast-components
|
||||
veo2-pro
|
||||
veo3
|
||||
veo3-fast
|
||||
veo3-pro
|
||||
veo3-pro-frames
|
||||
veo3-fast-frames
|
||||
veo3-frames
|
||||
|
||||
NOTE: veo3 does not support first and last frame to video generation.
|
||||
"""
|
||||
self.base_url = base_url.rstrip("/")
|
||||
self.api_key = api_key
|
||||
self.t2v_model = t2v_model
|
||||
self.ff2v_model = ff2v_model
|
||||
self.flf2v_model = flf2v_model
|
||||
|
||||
async def generate_single_video(
|
||||
self,
|
||||
prompt: str = "",
|
||||
reference_image_paths: List[Image.Image] = [],
|
||||
aspect_ratio: str = "16:9",
|
||||
**kwargs,
|
||||
) -> VideoOutput:
|
||||
progress = kwargs.get("progress")
|
||||
create_retries = _env_int("VIMAX_VIDEO_CREATE_RETRIES", 3)
|
||||
query_timeout_seconds = _env_float("VIMAX_VIDEO_QUERY_TIMEOUT_SECONDS", 600.0)
|
||||
request_timeout_seconds = _env_float("VIMAX_VIDEO_REQUEST_TIMEOUT_SECONDS", 60.0)
|
||||
poll_interval_seconds = _env_float("VIMAX_VIDEO_POLL_INTERVAL_SECONDS", 5.0)
|
||||
max_query_errors = _env_int("VIMAX_VIDEO_MAX_QUERY_ERRORS", 5)
|
||||
if len(reference_image_paths) == 0:
|
||||
model = self.t2v_model
|
||||
elif len(reference_image_paths) == 1:
|
||||
model = self.ff2v_model
|
||||
elif len(reference_image_paths) == 2:
|
||||
model = self.flf2v_model
|
||||
else:
|
||||
raise ValueError("The number of reference images must be no more than 2")
|
||||
|
||||
logging.info(f"Calling {model} to generate video...")
|
||||
|
||||
# 1. Create video generation task
|
||||
payload = {
|
||||
"prompt": prompt,
|
||||
"model": model,
|
||||
"images": [image_path_to_b64(image_path, mime=True) for image_path in reference_image_paths],
|
||||
"enhance_prompt": True,
|
||||
}
|
||||
# only veo3 supports aspect ratio setting
|
||||
if model.startswith("veo3"):
|
||||
payload["aspect_ratio"] = aspect_ratio
|
||||
|
||||
headers = {
|
||||
"Accept": "application/json",
|
||||
"Authorization": f"Bearer {self.api_key}",
|
||||
"Content-Type": "application/json",
|
||||
}
|
||||
|
||||
url = f"{self.base_url}/v1/video/create"
|
||||
task_id = None
|
||||
last_create_error = None
|
||||
timeout = aiohttp.ClientTimeout(total=request_timeout_seconds)
|
||||
for attempt in range(1, create_retries + 1):
|
||||
try:
|
||||
_emit_progress(progress, "video_create", f"Creating video generation task with {model}", {"model": model, "attempt": attempt, "max_attempts": create_retries})
|
||||
async with aiohttp.ClientSession(timeout=timeout) as session:
|
||||
async with session.post(url, headers=headers, json=payload) as response:
|
||||
response_payload = await response.json(content_type=None)
|
||||
logging.debug(f"Response: {response_payload}")
|
||||
if response.status >= 400:
|
||||
raise RuntimeError(f"Video create failed with HTTP {response.status}: {response_payload}")
|
||||
task_id = response_payload.get("id")
|
||||
if not task_id:
|
||||
raise RuntimeError(f"Video create response missing id: {response_payload}")
|
||||
logging.info(f"Video generation task created successfully. Task ID: {task_id}")
|
||||
_emit_progress(progress, "video_task_created", "Video generation task created", {"model": model, "task_id": task_id})
|
||||
break
|
||||
except Exception as e:
|
||||
last_create_error = e
|
||||
logging.error(f"Error occurred while creating video generation task: {e}.")
|
||||
_emit_progress(progress, "video_create_error", f"Video create attempt {attempt} failed", {"model": model, "attempt": attempt, "error": str(e)})
|
||||
if attempt < create_retries:
|
||||
await asyncio.sleep(1)
|
||||
if not task_id:
|
||||
raise RuntimeError(f"Video create failed after {create_retries} attempts: {last_create_error}")
|
||||
|
||||
|
||||
# 2. Query the video generation task until the video generation is completed
|
||||
headers = {
|
||||
'Accept': 'application/json',
|
||||
'Authorization': f'Bearer {self.api_key}',
|
||||
}
|
||||
|
||||
deadline = asyncio.get_running_loop().time() + query_timeout_seconds if query_timeout_seconds > 0 else None
|
||||
query_errors = 0
|
||||
last_status = None
|
||||
while deadline is None or asyncio.get_running_loop().time() < deadline:
|
||||
try:
|
||||
async with aiohttp.ClientSession(timeout=timeout) as session:
|
||||
async with session.get(f"{self.base_url}/v1/video/query?id={task_id}", headers=headers) as response:
|
||||
payload = await response.json(content_type=None)
|
||||
logging.debug(f"Response: {payload}")
|
||||
if response.status >= 400:
|
||||
raise RuntimeError(f"Video query failed with HTTP {response.status}: {payload}")
|
||||
status = payload.get("status")
|
||||
if not status:
|
||||
raise RuntimeError(f"Video query response missing status: {payload}")
|
||||
query_errors = 0
|
||||
except Exception as e:
|
||||
query_errors += 1
|
||||
logging.error(f"Error occurred while querying video generation task: {e}.")
|
||||
_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})
|
||||
if query_errors >= max_query_errors:
|
||||
raise RuntimeError(f"Video query failed {query_errors} times for task {task_id}: {e}")
|
||||
await asyncio.sleep(poll_interval_seconds)
|
||||
continue
|
||||
|
||||
if status == "completed":
|
||||
logging.info(f"Video generation completed successfully")
|
||||
video_url = payload.get("video_url")
|
||||
if not video_url:
|
||||
raise RuntimeError(f"Video task completed without video_url: {payload}")
|
||||
_emit_progress(progress, "video_completed", "Video generation completed", {"model": model, "task_id": task_id})
|
||||
return VideoOutput(fmt="url", ext="mp4", data=video_url)
|
||||
elif status == "failed":
|
||||
logging.error(f"Video generation failed: \n{payload}")
|
||||
raise RuntimeError(f"Video generation failed for task {task_id}: {payload}")
|
||||
else:
|
||||
logging.info(f"Video generation status: {status}, waiting 1 second...")
|
||||
last_status = status
|
||||
_emit_progress(progress, "video_status", f"Video generation status: {status}", {"model": model, "task_id": task_id, "status": status})
|
||||
await asyncio.sleep(poll_interval_seconds)
|
||||
continue
|
||||
raise RuntimeError(f"Video generation timed out after {query_timeout_seconds:g}s for task {task_id}; last_status={last_status}")
|
||||
Reference in New Issue
Block a user