AI Researcher and Video Studio implementation complete
This commit is contained in:
ajaysi
161
backend/services/wavespeed/generators/video/extension.py
Normal file
161
backend/services/wavespeed/generators/video/extension.py
Normal file
@@ -0,0 +1,161 @@
|
||||
"""
|
||||
Video extension operations.
|
||||
"""
|
||||
|
||||
import requests
|
||||
from typing import Optional, Callable
|
||||
from fastapi import HTTPException
|
||||
|
||||
from utils.logger_utils import get_service_logger
|
||||
from .base import VideoBase
|
||||
|
||||
logger = get_service_logger("wavespeed.generators.video.extension")
|
||||
|
||||
|
||||
class VideoExtension(VideoBase):
|
||||
"""Video extension operations."""
|
||||
|
||||
def extend_video(
|
||||
self,
|
||||
video: str, # Base64-encoded video or URL
|
||||
prompt: str,
|
||||
model: str = "wan-2.5", # "wan-2.5", "wan-2.2-spicy", or "seedance-1.5-pro"
|
||||
audio: Optional[str] = None, # Optional audio URL (WAN 2.5 only)
|
||||
negative_prompt: Optional[str] = None, # WAN 2.5 only
|
||||
resolution: str = "720p",
|
||||
duration: int = 5,
|
||||
enable_prompt_expansion: bool = False, # WAN 2.5 only
|
||||
generate_audio: bool = True, # Seedance 1.5 Pro only
|
||||
camera_fixed: bool = False, # Seedance 1.5 Pro only
|
||||
seed: Optional[int] = None,
|
||||
enable_sync_mode: bool = False,
|
||||
timeout: int = 300,
|
||||
progress_callback: Optional[Callable[[float, str], None]] = None,
|
||||
) -> bytes:
|
||||
"""
|
||||
Extend video duration using WAN 2.5, WAN 2.2 Spicy, or Seedance 1.5 Pro video-extend.
|
||||
|
||||
Args:
|
||||
video: Base64-encoded video data URI or public URL
|
||||
prompt: Text prompt describing how to extend the video
|
||||
model: Model to use ("wan-2.5", "wan-2.2-spicy", or "seedance-1.5-pro")
|
||||
audio: Optional audio URL to guide generation (WAN 2.5 only)
|
||||
negative_prompt: Optional negative prompt (WAN 2.5 only)
|
||||
resolution: Output resolution (varies by model)
|
||||
duration: Duration of extended video in seconds (varies by model)
|
||||
enable_prompt_expansion: Enable prompt optimizer (WAN 2.5 only)
|
||||
generate_audio: Generate audio for extended video (Seedance 1.5 Pro only)
|
||||
camera_fixed: Fix camera position (Seedance 1.5 Pro only)
|
||||
seed: Random seed for reproducibility (-1 for random)
|
||||
enable_sync_mode: If True, wait for result and return it directly
|
||||
timeout: Request timeout in seconds (default: 300)
|
||||
progress_callback: Optional callback function(progress: float, message: str) for progress updates
|
||||
|
||||
Returns:
|
||||
bytes: Extended video bytes
|
||||
|
||||
Raises:
|
||||
HTTPException: If the extension fails
|
||||
"""
|
||||
# Determine model path
|
||||
if model in ("wan-2.2-spicy", "wavespeed-ai/wan-2.2-spicy/video-extend"):
|
||||
model_path = "wavespeed-ai/wan-2.2-spicy/video-extend"
|
||||
elif model in ("seedance-1.5-pro", "bytedance/seedance-v1.5-pro/video-extend"):
|
||||
model_path = "bytedance/seedance-v1.5-pro/video-extend"
|
||||
else:
|
||||
# Default to WAN 2.5
|
||||
model_path = "alibaba/wan-2.5/video-extend"
|
||||
|
||||
url = f"{self.base_url}/{model_path}"
|
||||
|
||||
# Base payload (common to all models)
|
||||
payload = {
|
||||
"video": video,
|
||||
"prompt": prompt,
|
||||
"resolution": resolution,
|
||||
"duration": duration,
|
||||
}
|
||||
|
||||
# Model-specific parameters
|
||||
if model_path == "alibaba/wan-2.5/video-extend":
|
||||
# WAN 2.5 specific
|
||||
payload["enable_prompt_expansion"] = enable_prompt_expansion
|
||||
if audio:
|
||||
payload["audio"] = audio
|
||||
if negative_prompt:
|
||||
payload["negative_prompt"] = negative_prompt
|
||||
elif model_path == "bytedance/seedance-v1.5-pro/video-extend":
|
||||
# Seedance 1.5 Pro specific
|
||||
payload["generate_audio"] = generate_audio
|
||||
payload["camera_fixed"] = camera_fixed
|
||||
|
||||
# Seed (all models support it)
|
||||
if seed is not None:
|
||||
payload["seed"] = seed
|
||||
|
||||
logger.info(f"[WaveSpeed] Extending video via {url} (duration={duration}s, resolution={resolution})")
|
||||
|
||||
# Submit the task
|
||||
response = requests.post(url, headers=self._get_headers(), json=payload, timeout=timeout)
|
||||
|
||||
if response.status_code != 200:
|
||||
logger.error(f"[WaveSpeed] Video extend submission failed: {response.status_code} {response.text}")
|
||||
raise HTTPException(
|
||||
status_code=502,
|
||||
detail={
|
||||
"error": "WaveSpeed video extend submission failed",
|
||||
"status_code": response.status_code,
|
||||
"response": response.text,
|
||||
},
|
||||
)
|
||||
|
||||
response_json = response.json()
|
||||
data = response_json.get("data") or response_json
|
||||
prediction_id = data.get("id")
|
||||
|
||||
if not prediction_id:
|
||||
logger.error(f"[WaveSpeed] No prediction ID in video extend response: {response.text}")
|
||||
raise HTTPException(
|
||||
status_code=502,
|
||||
detail="WaveSpeed video extend response missing prediction id",
|
||||
)
|
||||
|
||||
logger.info(f"[WaveSpeed] Video extend task submitted: {prediction_id}")
|
||||
|
||||
# Poll for result
|
||||
result = self.polling.poll_until_complete(
|
||||
prediction_id,
|
||||
timeout_seconds=timeout,
|
||||
interval_seconds=2.0,
|
||||
progress_callback=progress_callback,
|
||||
)
|
||||
|
||||
outputs = result.get("outputs") or []
|
||||
if not outputs:
|
||||
raise HTTPException(status_code=502, detail="WaveSpeed video extend returned no outputs")
|
||||
|
||||
# Handle outputs - can be array of strings or array of objects
|
||||
video_url = None
|
||||
if isinstance(outputs[0], str):
|
||||
video_url = outputs[0]
|
||||
elif isinstance(outputs[0], dict):
|
||||
video_url = outputs[0].get("url") or outputs[0].get("video_url")
|
||||
|
||||
if not video_url:
|
||||
raise HTTPException(status_code=502, detail="WaveSpeed video extend output format not recognized")
|
||||
|
||||
# Download the extended video
|
||||
logger.info(f"[WaveSpeed] Downloading extended video from: {video_url}")
|
||||
video_response = requests.get(video_url, timeout=timeout)
|
||||
|
||||
if video_response.status_code != 200:
|
||||
logger.error(f"[WaveSpeed] Failed to download extended video: {video_response.status_code}")
|
||||
raise HTTPException(
|
||||
status_code=502,
|
||||
detail="Failed to download extended video from WaveSpeed",
|
||||
)
|
||||
|
||||
video_bytes = video_response.content
|
||||
logger.info(f"[WaveSpeed] Video extension completed successfully (size: {len(video_bytes)} bytes)")
|
||||
|
||||
return video_bytes
|
||||
Reference in New Issue
Block a user