aria-code 0.167.0__py3-none-win_amd64.whl

aria_code/video_workflow/engines/custom_model.py

@@ -0,0 +1,204 @@
+"""custom_model.py — 自研私有视频模型引擎适配器。
+
+专为用户后续自研、微调或私有化部署的视频生成大模型（如自建 DiT 架构、ComfyUI 服务端、
+vLLM / Triton 推理端点、Vertex AI Custom Prediction 或独立 GPU 云服务）打造。
+通过标准化的 HTTP/REST 接口规范与可配置的 Payload 模板，支持随时将底层生成驱动
+切换为自己的模型。
+"""
+
+from __future__ import annotations
+
+import base64
+import json
+import logging
+import os
+import time
+from pathlib import Path
+from typing import Any, Dict, Optional
+
+import requests
+
+from .base import BaseVideoEngine, VideoShotRequest, VideoShotResult
+from aria_code.artifacts import create_user_artifact
+
+logger = logging.getLogger(__name__)
+
+DEFAULT_CUSTOM_ENDPOINT = "http://localhost:8000/v1/video/generate"
+
+
+class CustomModelEngine(BaseVideoEngine):
+    """自研视频大模型适配引擎。"""
+
+    name: str = "custom"
+    display_name: str = "Custom Video Model (Self-Hosted / Private Cloud)"
+
+    def __init__(
+        self,
+        endpoint_url: Optional[str] = None,
+        auth_token: Optional[str] = None,
+        model_name: Optional[str] = None,
+    ):
+        self.endpoint_url = endpoint_url or os.getenv("ARIA_CUSTOM_VIDEO_ENDPOINT", DEFAULT_CUSTOM_ENDPOINT)
+        self.auth_token = auth_token or os.getenv("ARIA_CUSTOM_VIDEO_TOKEN", "")
+        self.model_name = model_name or os.getenv("ARIA_CUSTOM_VIDEO_MODEL", "aria-custom-dit-v1")
+
+    def health_check(self) -> Dict[str, Any]:
+        """检查自研模型端点的联通性。"""
+        is_reachable = False
+        error_msg = ""
+        try:
+            health_url = self.endpoint_url.rsplit("/", 1)[0] + "/health"
+            resp = requests.get(health_url, timeout=3)
+            is_reachable = resp.status_code == 200
+        except Exception as exc:
+            error_msg = str(exc)
+
+        return {
+            "engine": self.name,
+            "display_name": self.display_name,
+            "endpoint_url": self.endpoint_url,
+            "model_name": self.model_name,
+            "has_auth_token": bool(self.auth_token),
+            "reachable": is_reachable,
+            "status_message": "OK" if is_reachable else f"Not reachable ({error_msg or 'no health endpoint'})",
+        }
+
+    def generate_shot(self, request: VideoShotRequest) -> VideoShotResult:
+        """向自研模型服务提交生成请求，并保存结果。"""
+        headers = {"Content-Type": "application/json"}
+        if self.auth_token:
+            headers["Authorization"] = f"Bearer {self.auth_token}"
+
+        # 构造自研模型标准请求负载
+        payload: Dict[str, Any] = {
+            "model": self.model_name,
+            "prompt": request.prompt,
+            "negative_prompt": request.negative_prompt or "blurry, low quality, artifacts, distorted",
+            "duration": request.duration_seconds,
+            "aspect_ratio": request.aspect_ratio,
+            "cfg_scale": request.cfg_scale,
+        }
+
+        if request.seed is not None:
+            payload["seed"] = request.seed
+        if request.camera_motion:
+            payload["camera_motion"] = request.camera_motion
+
+        # 处理首帧图片输入 (图生视频 / I2V)
+        if request.image_path and Path(request.image_path).exists():
+            img_bytes = Path(request.image_path).read_bytes()
+            payload["image_base64"] = base64.b64encode(img_bytes).decode("ascii")
+
+        # 混入自定义扩展参数 (例如 lora_weights, sampling_steps, motion_scale)
+        if request.custom_parameters:
+            payload.update(request.custom_parameters)
+
+        logger.info(f"[CustomModelEngine] Dispatching to {self.endpoint_url} (model={self.model_name})...")
+
+        try:
+            resp = requests.post(self.endpoint_url, headers=headers, json=payload, timeout=120)
+            if resp.status_code != 200:
+                return VideoShotResult(
+                    success=False,
+                    engine_name=self.name,
+                    model_version=self.model_name,
+                    error=f"Custom model endpoint returned HTTP {resp.status_code}: {resp.text[:300]}",
+                )
+
+            data = resp.json()
+
+            # 1. 如果端点直接返回 base64 视频流
+            video_b64 = data.get("video_base64")
+            if video_b64:
+                artifact = create_user_artifact("video", "custom_model", "clip", ".mp4")
+                artifact.path.write_bytes(base64.b64decode(video_b64))
+                return VideoShotResult(
+                    success=True,
+                    path=str(artifact.path),
+                    engine_name=self.name,
+                    model_version=self.model_name,
+                    duration=request.duration_seconds,
+                    aspect_ratio=request.aspect_ratio,
+                    metadata=data.get("metadata", {}),
+                )
+
+            # 2. 如果端点返回可下载视频 URL
+            video_url = data.get("video_url")
+            if video_url:
+                dl_resp = requests.get(video_url, timeout=120)
+                if dl_resp.status_code == 200:
+                    artifact = create_user_artifact("video", "custom_model", "clip", ".mp4")
+                    artifact.path.write_bytes(dl_resp.content)
+                    return VideoShotResult(
+                        success=True,
+                        path=str(artifact.path),
+                        video_url=video_url,
+                        engine_name=self.name,
+                        model_version=self.model_name,
+                        duration=request.duration_seconds,
+                        aspect_ratio=request.aspect_ratio,
+                        metadata=data.get("metadata", {}),
+                    )
+
+            # 3. 如果是异步任务 (task_id / poll_url)
+            task_id = data.get("task_id")
+            poll_url = data.get("poll_url") or f"{self.endpoint_url.rsplit('/', 1)[0]}/tasks/{task_id}"
+            if task_id:
+                return self._poll_task(poll_url, headers, request.duration_seconds, request.aspect_ratio)
+
+            return VideoShotResult(
+                success=False,
+                engine_name=self.name,
+                model_version=self.model_name,
+                error="Custom endpoint response contained neither video_base64, video_url nor task_id.",
+            )
+        except Exception as exc:
+            logger.exception(f"[CustomModelEngine] Error calling custom endpoint: {exc}")
+            return VideoShotResult(
+                success=False,
+                engine_name=self.name,
+                model_version=self.model_name,
+                error=str(exc),
+            )
+
+    def _poll_task(self, poll_url: str, headers: Dict[str, str], duration: int, aspect: str, max_wait: float = 300.0) -> VideoShotResult:
+        """异步轮询任务结果。"""
+        start = time.time()
+        while time.time() - start < max_wait:
+            time.sleep(3.0)
+            try:
+                res = requests.get(poll_url, headers=headers, timeout=20)
+                if res.status_code == 200:
+                    pdata = res.json()
+                    status = pdata.get("status")
+                    if status in ("succeed", "completed", "success"):
+                        video_url = pdata.get("video_url")
+                        if video_url:
+                            v_resp = requests.get(video_url, timeout=120)
+                            artifact = create_user_artifact("video", "custom_model", "clip", ".mp4")
+                            artifact.path.write_bytes(v_resp.content)
+                            return VideoShotResult(
+                                success=True,
+                                path=str(artifact.path),
+                                video_url=video_url,
+                                engine_name=self.name,
+                                model_version=self.model_name,
+                                duration=duration,
+                                aspect_ratio=aspect,
+                            )
+                    elif status in ("failed", "error"):
+                        return VideoShotResult(
+                            success=False,
+                            engine_name=self.name,
+                            model_version=self.model_name,
+                            error=pdata.get("error", "Task execution failed on custom model."),
+                        )
+            except Exception as exc:
+                logger.debug(f"[CustomModelEngine] Poll error: {exc}")
+
+        return VideoShotResult(
+            success=False,
+            engine_name=self.name,
+            model_version=self.model_name,
+            error="Custom model task polling timed out.",
+        )

aria_code/video_workflow/engines/google_veo.py

@@ -0,0 +1,94 @@
+"""google_veo.py — Google Cloud Vertex AI & Gemini API Veo 视频引擎实现。"""
+
+from __future__ import annotations
+
+import logging
+import os
+from typing import Any, Dict
+
+from .base import BaseVideoEngine, VideoShotRequest, VideoShotResult
+from aria_code.providers.google_video_provider import has_google_video_config, generate_video_google, DEFAULT_VEO_MODEL
+
+logger = logging.getLogger(__name__)
+
+
+class GoogleCloudVeoEngine(BaseVideoEngine):
+    """Google Cloud Veo 视频生成引擎（当前主要生产驱动）。"""
+
+    name: str = "google"
+    display_name: str = "Google Cloud Veo (Vertex AI / Gemini API)"
+
+    def __init__(self, model: str = DEFAULT_VEO_MODEL):
+        self.model = model
+
+    def health_check(self) -> Dict[str, Any]:
+        """检查 Google Cloud 凭证与配置状态。"""
+        vertex_project = os.getenv("VERTEX_PROJECT") or os.getenv("GOOGLE_CLOUD_PROJECT", "")
+        api_key = os.getenv("GEMINI_API_KEY", "") or os.getenv("GOOGLE_API_KEY", "")
+        configured = has_google_video_config()
+
+        mode = "Vertex AI (ADC)" if vertex_project else ("Gemini API Key" if api_key else "None")
+        return {
+            "engine": self.name,
+            "display_name": self.display_name,
+            "configured": configured,
+            "auth_mode": mode,
+            "model": self.model,
+            "project_id": vertex_project if vertex_project else "(using Gemini API Key)",
+        }
+
+    def generate_shot(self, request: VideoShotRequest) -> VideoShotResult:
+        """调用 Google Cloud Veo 生成单个视频片段。"""
+        logger.info(f"[GoogleVeoEngine] Submitting shot: {request.prompt[:60]}...")
+
+        # 丰富 Prompt：若有镜头运镜参数则显式融入 Prompt
+        enhanced_prompt = request.prompt
+        if request.camera_motion:
+            motions = []
+            if request.camera_motion.get("zoom", 0) > 0:
+                motions.append("camera slow zoom in")
+            elif request.camera_motion.get("zoom", 0) < 0:
+                motions.append("camera slow zoom out")
+            if request.camera_motion.get("pan", 0) > 0:
+                motions.append("camera pan right")
+            elif request.camera_motion.get("pan", 0) < 0:
+                motions.append("camera pan left")
+            if request.camera_motion.get("tilt", 0) > 0:
+                motions.append("camera tilt up")
+            elif request.camera_motion.get("tilt", 0) < 0:
+                motions.append("camera tilt down")
+            if motions:
+                enhanced_prompt = f"{enhanced_prompt}. {', '.join(motions)}."
+
+        try:
+            res = generate_video_google(
+                prompt=enhanced_prompt,
+                image_path=request.image_path,
+                model=self.model,
+                aspect_ratio=request.aspect_ratio,
+                duration_seconds=request.duration_seconds,
+            )
+            if res.get("success") and res.get("path"):
+                return VideoShotResult(
+                    success=True,
+                    path=res["path"],
+                    engine_name=self.name,
+                    model_version=res.get("model", self.model),
+                    duration=float(res.get("duration", request.duration_seconds)),
+                    aspect_ratio=res.get("aspect_ratio", request.aspect_ratio),
+                    metadata={"provider": res.get("provider", "Google Cloud")},
+                )
+            return VideoShotResult(
+                success=False,
+                engine_name=self.name,
+                model_version=self.model,
+                error=res.get("error", "Google Veo failed to produce video."),
+            )
+        except Exception as exc:
+            logger.exception(f"[GoogleVeoEngine] Generation exception: {exc}")
+            return VideoShotResult(
+                success=False,
+                engine_name=self.name,
+                model_version=self.model,
+                error=str(exc),
+            )

aria_code/video_workflow/engines/local_engine.py

@@ -0,0 +1,71 @@
+"""local_engine.py — 本地离线高质感运镜生成引擎（免网络、零成本兜底）。"""
+
+from __future__ import annotations
+
+import logging
+from pathlib import Path
+from typing import Any, Dict
+
+from .base import BaseVideoEngine, VideoShotRequest, VideoShotResult
+from aria_code.video_editor import create_ken_burns_clip
+
+logger = logging.getLogger(__name__)
+
+
+class LocalFallbackEngine(BaseVideoEngine):
+    """本地运镜引擎，当云端与私有模型不可用或用户需要零成本预览时使用。"""
+
+    name: str = "local"
+    display_name: str = "Local Ken-Burns Motion Engine (Offline)"
+
+    def health_check(self) -> Dict[str, Any]:
+        import shutil
+        has_ffmpeg = shutil.which("ffmpeg") is not None
+        return {
+            "engine": self.name,
+            "display_name": self.display_name,
+            "configured": has_ffmpeg,
+            "ffmpeg_available": has_ffmpeg,
+            "status_message": "Ready (Offline)" if has_ffmpeg else "ffmpeg not found in PATH",
+        }
+
+    def generate_shot(self, request: VideoShotRequest) -> VideoShotResult:
+        """从图片素材本地渲染 Ken Burns 运镜视频。"""
+        if not request.image_path or not Path(request.image_path).exists():
+            return VideoShotResult(
+                success=False,
+                engine_name=self.name,
+                error="Local engine requires a valid source image (image_path).",
+            )
+
+        effect = "zoom_in"
+        if request.camera_motion:
+            if request.camera_motion.get("pan", 0) > 0:
+                effect = "pan_right"
+            elif request.camera_motion.get("pan", 0) < 0:
+                effect = "pan_left"
+            elif request.camera_motion.get("zoom", 0) < 0:
+                effect = "zoom_out"
+
+        res = create_ken_burns_clip(
+            request.image_path,
+            duration=float(request.duration_seconds),
+            effect=effect,
+            aspect_ratio=request.aspect_ratio,
+        )
+
+        if res.get("success") and res.get("clip_path"):
+            return VideoShotResult(
+                success=True,
+                path=res["clip_path"],
+                engine_name=self.name,
+                model_version="ken_burns_v2",
+                duration=float(request.duration_seconds),
+                aspect_ratio=request.aspect_ratio,
+            )
+
+        return VideoShotResult(
+            success=False,
+            engine_name=self.name,
+            error=res.get("error", "Local Ken Burns generation failed."),
+        )

aria_code/video_workflow/engines/registry.py

@@ -0,0 +1,98 @@
+"""registry.py — 视频生成引擎注册表与生命周期管理器。
+
+支持全局统一管理 Google Cloud Veo、自研私有模型 (Custom) 与本地引擎 (Local)，
+支持通过环境变量、代码显式配置或 MCP 接口随时切换当前生效的主视频引擎。
+"""
+
+from __future__ import annotations
+
+import logging
+import os
+from typing import Any, Dict, List, Optional, Type
+
+from .base import BaseVideoEngine
+from .google_veo import GoogleCloudVeoEngine
+from .custom_model import CustomModelEngine
+from .local_engine import LocalFallbackEngine
+
+logger = logging.getLogger(__name__)
+
+
+class VideoEngineRegistry:
+    """引擎注册表单例管理器。"""
+
+    _instance: Optional[VideoEngineRegistry] = None
+
+    def __init__(self):
+        self._engines: Dict[str, BaseVideoEngine] = {}
+        # 默认内置三大引擎
+        self.register_engine("google", GoogleCloudVeoEngine())
+        self.register_engine("custom", CustomModelEngine())
+        self.register_engine("local", LocalFallbackEngine())
+        self._active_engine_name: Optional[str] = None
+
+    @classmethod
+    def get_instance(cls) -> VideoEngineRegistry:
+        if cls._instance is None:
+            cls._instance = cls()
+        return cls._instance
+
+    @classmethod
+    def reset_instance(cls) -> None:
+        """Drop the singleton so the next caller builds a fresh registry.
+
+        The registry is process-wide and mutable: registering an engine or
+        setting the active one changes it for everything afterwards. In a test
+        run that leaks — configuring a custom engine in one test left another
+        patching "google" while the director resolved the active engine and
+        called something else, which failed only when both ran together.
+        """
+        cls._instance = None
+
+    def register_engine(self, name: str, engine: BaseVideoEngine) -> None:
+        """注册或更新一个视频引擎实例。"""
+        self._engines[name.lower()] = engine
+        logger.debug(f"[VideoEngineRegistry] Registered engine '{name}': {engine.display_name}")
+
+    def set_active_engine(self, name: str) -> None:
+        """设置当前全局默认视频引擎。"""
+        name_lower = name.lower()
+        if name_lower not in self._engines:
+            raise ValueError(f"Unknown engine '{name}'. Available: {list(self._engines.keys())}")
+        self._active_engine_name = name_lower
+        logger.info(f"[VideoEngineRegistry] Active video engine set to: {name_lower}")
+
+    def get_active_engine_name(self) -> str:
+        """获取当前活动引擎名称。优先按环境变量 ARIA_VIDEO_ENGINE，其次按配置，最后默认 google。"""
+        if self._active_engine_name:
+            return self._active_engine_name
+        env_choice = os.getenv("ARIA_VIDEO_ENGINE", "").strip().lower()
+        if env_choice in self._engines:
+            return env_choice
+        # 若自研模型端点被明确指定，可优先考虑
+        if os.getenv("ARIA_CUSTOM_VIDEO_ENDPOINT"):
+            return "custom"
+        return "google"
+
+    def get_engine(self, name: Optional[str] = None) -> BaseVideoEngine:
+        """获取指定或当前生效的视频引擎实例。"""
+        target = (name or self.get_active_engine_name()).lower()
+        if target not in self._engines:
+            logger.warning(f"Engine '{target}' not found, falling back to 'local'.")
+            target = "local"
+        return self._engines[target]
+
+    def list_engines_status(self) -> List[Dict[str, Any]]:
+        """获取所有已注册引擎的运行与配置健康度报告。"""
+        active_name = self.get_active_engine_name()
+        status_list = []
+        for name, engine in self._engines.items():
+            check = engine.health_check()
+            check["is_active"] = (name == active_name)
+            status_list.append(check)
+        return status_list
+
+
+def get_video_engine(name: Optional[str] = None) -> BaseVideoEngine:
+    """快捷获取当前或指定的视频引擎。"""
+    return VideoEngineRegistry.get_instance().get_engine(name)