diff --git a/tests/tools/test_cogvideo_i2v_variant.py b/tests/tools/test_cogvideo_i2v_variant.py new file mode 100644 index 00000000..8efa7fca --- /dev/null +++ b/tests/tools/test_cogvideo_i2v_variant.py @@ -0,0 +1,78 @@ +"""Regression for cogvideo-2b i2v mismatch (REVIEW §8 #4). + +COGVIDEO_VARIANTS declares the 2B variant i2v=False (t2v-only), but +cogvideo_video advertised image_to_video + reference_image unconditionally +and never consulted the variant flag — so an image_to_video brief against +cogvideo-2b reached the diffusion pipeline and failed opaquely. + +This pins: +- is_operation_available() derives capability from the variant table + (default 5b: t2v + i2v both True), not an unconditional True. +- execute() fails fast with a clear error when the caller's chosen variant + lacks the requested mode (2b + image_to_video). +""" + +from __future__ import annotations + +import pytest + +from tools.base_tool import ToolStatus +from tools.video._shared import COGVIDEO_VARIANTS +from tools.video.cogvideo_video import CogVideoVideo + + +def test_2b_variant_is_t2v_only(): + """The premise: the 2B variant really does declare i2v=False.""" + assert COGVIDEO_VARIANTS["cogvideo-2b"]["i2v"] is False + assert COGVIDEO_VARIANTS["cogvideo-2b"]["t2v"] is True + + +def test_5b_variant_supports_i2v(): + assert COGVIDEO_VARIANTS["cogvideo-5b"]["i2v"] is True + + +def test_is_operation_available_reflects_default_variant(): + tool = CogVideoVideo() + # Default variant (5b) supports both modes. + assert tool.is_operation_available("text_to_video") is True + assert tool.is_operation_available("image_to_video") is True + + +def test_execute_rejects_i2v_for_2b_variant(monkeypatch): + """Selecting the 2B variant for image_to_video fails fast, not at the pipeline.""" + monkeypatch.setattr( + CogVideoVideo, "get_status", lambda self: ToolStatus.AVAILABLE + ) + # Guard against the local-generation path actually running. + monkeypatch.setattr( + "tools.video.cogvideo_video.generate_local_video", + lambda **kw: pytest.fail("generate_local_video must not run for an unsupported variant"), + ) + + result = CogVideoVideo().execute( + {"prompt": "x", "operation": "image_to_video", "model_variant": "cogvideo-2b"} + ) + assert result.success is False + assert "cogvideo-2b" in result.error + assert "image_to_video" in result.error + + +def test_execute_allows_i2v_for_5b_variant(monkeypatch): + """The default 5B variant still routes to generation for image_to_video.""" + from tools.base_tool import ToolResult + + monkeypatch.setattr( + CogVideoVideo, "get_status", lambda self: ToolStatus.AVAILABLE + ) + + sentinel = ToolResult(success=True, data={"output": "out.mp4"}) + + def fake_generate(**kw): + return sentinel + + monkeypatch.setattr("tools.video.cogvideo_video.generate_local_video", fake_generate) + + result = CogVideoVideo().execute( + {"prompt": "x", "operation": "image_to_video", "model_variant": "cogvideo-5b"} + ) + assert result.success is True diff --git a/tools/video/cogvideo_video.py b/tools/video/cogvideo_video.py index 9076fbd8..879cff10 100644 --- a/tools/video/cogvideo_video.py +++ b/tools/video/cogvideo_video.py @@ -77,6 +77,31 @@ class CogVideoVideo(BaseTool): def get_status(self) -> ToolStatus: return local_generation_status() + DEFAULT_VARIANT = "cogvideo-5b" + + def _variant_for(self, inputs: dict[str, object]) -> dict[str, object]: + """Resolve the active variant dict, falling back to the default.""" + return COGVIDEO_VARIANTS.get( + inputs.get("model_variant", self.DEFAULT_VARIANT), + COGVIDEO_VARIANTS[self.DEFAULT_VARIANT], + ) + + def is_operation_available(self, operation: str) -> bool: + """Capability is variant-driven, not unconditional. + + cogvideo_video advertised image_to_video unconditionally, but the 2B + variant declares i2v=False (it is t2v-only). The selector calls this + without inputs, so we report capability for the DEFAULT variant + (cogvideo-5b: t2v=True, i2v=True) — execute() re-checks against the + caller's chosen variant and fails fast if it lacks the requested mode. + """ + variant = COGVIDEO_VARIANTS[self.DEFAULT_VARIANT] + if operation == "image_to_video": + return bool(variant.get("i2v")) + if operation == "text_to_video": + return bool(variant.get("t2v")) + return True + def estimate_cost(self, inputs: dict[str, object]) -> float: return 0.0 @@ -87,6 +112,22 @@ class CogVideoVideo(BaseTool): def execute(self, inputs: dict[str, object]) -> ToolResult: if self.get_status() != ToolStatus.AVAILABLE: return ToolResult(success=False, error="CogVideo local generation is unavailable. " + self.install_instructions) + + # Consult the variant flag the selector cannot see: the 2B variant is + # t2v-only (i2v=False), so an image_to_video brief against it would + # otherwise reach the diffusion pipeline and fail opaquely. + operation = inputs.get("operation", "text_to_video") + variant = self._variant_for(inputs) + if operation == "image_to_video" and not variant.get("i2v"): + chosen = inputs.get("model_variant", self.DEFAULT_VARIANT) + return ToolResult( + success=False, + error=( + f"CogVideo variant '{chosen}' does not support image_to_video " + f"(i2v=False). Use the cogvideo-5b variant or a text_to_video operation." + ), + ) + start = time.time() try: result = generate_local_video(tool_name=self.name, variants=COGVIDEO_VARIANTS, default_variant="cogvideo-5b", inputs=inputs)