Fix long camera inference decode and scale guards

Author: Junchao-cs

configs/examples/wan22_ti2v_5b/infer_stage2_sgf_camera_length.yaml

@@ -87,6 +87,8 @@ validation:
   sample_count: 1
   selection_seed: 42
   noise_seed: 42
+  max_rel_translation: null
+  max_camera_abs: null
   passes:
     - {name: model_self_forcing_nfe4, weights: model, mode: autoregressive, solver: self_forcing, num_inference_steps: 4}
 checkpoint:

docs/backends/wan22-ti2v-5b.md

@@ -147,10 +147,13 @@ torchrun --standalone --nproc-per-node=1 -m solarwm infer \
   --set runtime.output_dir="$SOLAR_OUTPUT_ROOT/wan5-stage2-sgf-infer"
 ```
 
-The Stage2 command loads the checkpoint directory as one model and uses the
-longest camera-backed horizon available for each selected test sample. The
-horizon is rounded down to complete three-latent chunks; at the configured
-16 fps, a 960-frame camera track produces 240 latents and 957 decoded frames.
+The Stage2 command loads the checkpoint directory as one model, resolves its
+weight role from `release-manifest.json` (`ema` for the published checkpoint),
+and uses the longest camera-backed horizon available for each selected test
+sample. The horizon is rounded down to complete three-latent chunks; at the
+configured 16 fps, a 960-frame camera track produces 240 latents and 957
+decoded frames. Outputs longer than 60 latents are VAE-decoded as consecutive
+60-latent tiles with one continuous temporal cache.
 Override `data.test_index` to select a different raw-WDS test index. Outputs
 are written below `runtime.output_dir/generation`.
 

skills/solar-training-ops/SKILL.md

@@ -221,6 +221,9 @@ the weight selection declared by the chosen route before model allocation.
 The Wan5B Stage2 camera-length inference route accepts the released checkpoint
 directory as one model and does not expose LIVE/EMA selection. Accept a saved
 checkpoint only after its completion marker and declared members are durable.
+Resolve the model role from `release-manifest.json` and retain that resolved
+role in output provenance. Decode horizons longer than 60 latents as consecutive
+cached VAE tiles, clearing the temporal cache only before and after the sequence.
 
 ### 7. Keep inference and validation aligned
 

src/solarwm/backends/wan22/runtime/components.py

@@ -218,6 +218,48 @@ def decode(self, latents_btchw: Any, *, use_cache: bool = False) -> Any:
             output.append(decoded.float().clamp_(-1, 1).squeeze(0))
         return torch.stack(output, dim=0).permute(0, 2, 1, 3, 4)
 
+    def decode_streaming(
+        self,
+        latents_btchw: Any,
+        *,
+        chunk_latent_frames: int = 60,
+    ) -> Any:
+        """Decode consecutive temporal tiles with one continuous VAE cache."""
+
+        import torch
+
+        if latents_btchw.ndim != 5 or int(latents_btchw.shape[1]) <= 0:
+            raise BackendContractError("Wan streaming VAE decode requires non-empty BTCHW latents")
+        if chunk_latent_frames <= 0:
+            raise BackendContractError("Wan streaming VAE chunk size must be positive")
+        clear_cache = getattr(self.module, "clear_cache", None)
+        cached_decode = getattr(self.module, "cached_decode", None)
+        if not callable(clear_cache) or not callable(cached_decode):
+            raise BackendContractError(
+                "Wan streaming VAE decode requires cache-aware decoder methods"
+            )
+
+        clips = latents_btchw.permute(0, 2, 1, 3, 4)
+        output = []
+        for clip in clips:
+            decoded_chunks = []
+            clear_cache()
+            try:
+                for start in range(0, int(clip.shape[1]), chunk_latent_frames):
+                    chunk = clip[:, start : start + chunk_latent_frames].contiguous()
+                    with torch.autocast(device_type=chunk.device.type, dtype=chunk.dtype):
+                        decoded = cached_decode(chunk.unsqueeze(0), self._scale(chunk))
+                    if not bool(torch.isfinite(decoded).all().item()):
+                        raise BackendContractError(
+                            "Wan streaming VAE decode produced non-finite pixels"
+                        )
+                    decoded_chunks.append(decoded.float().clamp_(-1, 1).squeeze(0).cpu())
+                    del decoded
+            finally:
+                clear_cache()
+            output.append(torch.cat(decoded_chunks, dim=1))
+        return torch.stack(output, dim=0).permute(0, 2, 1, 3, 4)
+
 
 class WanA14BVAE:
     """Official deterministic 16-channel Wan2.1 VAE used by I2V-A14B."""

src/solarwm/backends/wan22/runtime/stage2.py

@@ -9,6 +9,7 @@
 
 import copy
 import gc
+import json
 import math
 import os
 import random
@@ -61,6 +62,7 @@
 
 _CHECKPOINT_FORMAT = "solarwm.wan22-stage2-sgf.v1"
 _TORCHRUN_OWNER_ENV = "SOLARWM_TORCHRUN_LIFECYCLE_OWNER"
+_STREAMING_VAE_LATENT_CHUNK = 60
 
 
 class Stage2GenerationRunner(Protocol):
@@ -830,9 +832,19 @@ def _stage2_generated_sample(
                 generation_pass.rollout_latent_frames,
             )
         )
-        decoded = provider.vae.decode(
-            latents[:, :output_latent_frames].contiguous(), use_cache=False
-        )
+        output_latents = latents[:, :output_latent_frames].contiguous()
+        if output_latent_frames > _STREAMING_VAE_LATENT_CHUNK:
+            decoded = provider.vae.decode_streaming(
+                output_latents,
+                chunk_latent_frames=_STREAMING_VAE_LATENT_CHUNK,
+            )
+            vae_decode = {
+                "mode": "continuous_cached_tiles",
+                "chunk_latent_frames": _STREAMING_VAE_LATENT_CHUNK,
+            }
+        else:
+            decoded = provider.vae.decode(output_latents, use_cache=False)
+            vae_decode = {"mode": "direct", "chunk_latent_frames": output_latent_frames}
     finite_fraction = float(torch.isfinite(decoded).float().mean().item())
     if finite_fraction != 1.0:
         raise BackendContractError("Stage2 VAE decode produced non-finite pixels")
@@ -868,6 +880,10 @@ def _stage2_generated_sample(
             "camera_translation_transform": str(
                 provider.config["model"]["camera_translation_transform"]
             ),
+            "resolved_weights_role": str(
+                getattr(provider, "_model_weight_role", generation_pass.weights)
+            ),
+            "vae_decode": vae_decode,
         },
     )
 
@@ -1040,6 +1056,38 @@ def _verified_stage2_inference_checkpoint(
     return requested, verified.manifest_digest, int(verified.step)
 
 
+def _published_default_weight_role(source: str | Path) -> str:
+    """Resolve the release-selected live/EMA role without exposing a CLI choice."""
+
+    requested = Path(source).expanduser().resolve()
+    root = requested if requested.is_dir() else requested.parent
+    manifest_path = root / "release-manifest.json"
+    try:
+        manifest = json.loads(manifest_path.read_text(encoding="utf-8"))
+    except Exception as exc:
+        raise BackendContractError(
+            f"Stage2 camera-length checkpoint lacks a readable release manifest: {exc}"
+        ) from exc
+    if not isinstance(manifest, Mapping):
+        raise BackendContractError("Stage2 release manifest must be a mapping")
+    load = manifest.get("load")
+    identity = manifest.get("identity")
+    model = identity.get("model") if isinstance(identity, Mapping) else None
+    if (
+        manifest.get("schema") != "solarwm.public-weight-manifest.v1"
+        or not isinstance(load, Mapping)
+        or load.get("format") != "solarwm_wan_stage2_transaction_v1"
+        or load.get("entrypoint") != "."
+        or not isinstance(model, Mapping)
+    ):
+        raise BackendContractError("Stage2 release manifest contract differs")
+    role = str(load.get("default_weights", "")).strip().lower()
+    available = str(model.get("weight_role", "")).strip().lower().split("+")
+    if role not in {"live", "ema"} or role not in available:
+        raise BackendContractError("Stage2 release default weight role is invalid")
+    return role
+
+
 def _plain(value: Any) -> Any:
     if isinstance(value, Mapping):
         return {str(key): _plain(item) for key, item in value.items()}
@@ -2236,6 +2284,11 @@ def __init__(self, values: Mapping[str, Any], generation_plan: Any) -> None:
                     values,
                     str(values["checkpoint"]["path"]),
                 )
+                self._model_weight_role = (
+                    _published_default_weight_role(values["checkpoint"]["path"])
+                    if self._direct_model
+                    else None
+                )
                 self.checkpoint_id = f"manifest:{checkpoint_manifest_id}"
                 layout = WanAssetLayout.from_config(values)
                 self.diffusion = build_diffusion_architecture(values)
@@ -2253,12 +2306,12 @@ def _root(module: Any) -> Any:
 
             def weight_id(self, role: str) -> str:
                 if role == "model" and self._direct_model:
-                    return self.checkpoint_id
+                    return f"{self.checkpoint_id}#release-default:{self._model_weight_role}"
                 return super().weight_id(role)
 
             def _checkpoint_state_field(self, role: str) -> str:
                 if role == "model" and self._direct_model:
-                    return "generator_ema"
+                    return super()._checkpoint_state_field(str(self._model_weight_role))
                 return super()._checkpoint_state_field(role)
 
             def allocate_kv_cache(

tests/backends/wan22/test_configs.py

@@ -157,6 +157,8 @@ def test_stage2_camera_length_inference_selects_one_direct_model() -> None:
     assert config["checkpoint"] == {"path": "/path/to/SolarWM-models/SolarWM-5B-sgf-stage2-81f"}
     assert config["inference"] == {"source": "validation", "length": "camera"}
     assert config["data"]["random_start"] is False
+    assert config["validation"]["max_rel_translation"] is None
+    assert config["validation"]["max_camera_abs"] is None
     assert [item["weights"] for item in config["validation"]["passes"]] == ["model"]
 
 

tests/backends/wan22/test_stage2_runtime.py

@@ -11,6 +11,7 @@
 from solarwm.backends.wan22.runtime.stage2 import (
     RoleCheckpointReceipt,
     Wan5BStage2Runtime,
+    _published_default_weight_role,
     _stage2_generated_sample,
     _stage2_initialization_receipt,
     _stage2_self_forcing_latents,
@@ -265,6 +266,97 @@ def test_stage2_generated_sample_preserves_configured_denoising_steps(
     assert generated.provenance["denoising_step_list"] == sentinel
 
 
+def test_stage2_long_generation_uses_streaming_vae_tiles(
+    monkeypatch: pytest.MonkeyPatch,
+) -> None:
+    torch = pytest.importorskip("torch")
+    from solarwm.backends.wan22.runtime import inference, stage2
+
+    calls: list[int] = []
+
+    class _VAE:
+        @staticmethod
+        def decode_streaming(value: object, *, chunk_latent_frames: int) -> object:
+            calls.extend((int(value.shape[1]), chunk_latent_frames))
+            return torch.zeros((1, 957, 3, 1, 1))
+
+        @staticmethod
+        def decode(*_args: object, **_kwargs: object) -> object:
+            raise AssertionError("long Stage2 generation must not use one-shot VAE decode")
+
+    monkeypatch.setattr(
+        stage2,
+        "_stage2_self_forcing_latents",
+        lambda *_: (
+            torch.zeros((1, 240, 1, 1, 1)),
+            {"timesteps": [1000, 750, 500, 250]},
+        ),
+    )
+    monkeypatch.setattr(inference, "_encode_compare_mp4", lambda *_args, **_kwargs: b"compare")
+    provider = SimpleNamespace(
+        device=torch.device("cpu"),
+        config={
+            "data": {"fps": 16.0},
+            "model": {"camera_translation_transform": "linear"},
+            "train": {"denoising_step_list": [1000, 750, 500, 250]},
+        },
+        _conditions=lambda *_args, **_kwargs: (
+            torch.zeros((1, 1, 1, 1, 1)),
+            {},
+            {},
+            None,
+        ),
+        vae=_VAE(),
+        video_encoder=lambda *_args, **_kwargs: b"video",
+        _prepared={0: object()},
+        _model_weight_role="ema",
+    )
+    case = SimpleNamespace(
+        slot=0,
+        noise_seed=42,
+        metadata={
+            "generation_pass": {
+                "name": "model_self_forcing_nfe4",
+                "weights": "model",
+                "mode": "autoregressive",
+                "solver": "self_forcing",
+                "num_inference_steps": 4,
+                "rollout_latent_frames": 240,
+            }
+        },
+    )
+
+    generated = _stage2_generated_sample(provider, case, weights_id="release#ema")
+
+    assert calls == [240, 60]
+    assert generated.shape == (1, 957, 3, 1, 1)
+    assert generated.provenance["resolved_weights_role"] == "ema"
+    assert generated.provenance["vae_decode"] == {
+        "mode": "continuous_cached_tiles",
+        "chunk_latent_frames": 60,
+    }
+
+
+def test_stage2_camera_length_resolves_release_default_weights(tmp_path: Path) -> None:
+    manifest = {
+        "schema": "solarwm.public-weight-manifest.v1",
+        "identity": {"model": {"weight_role": "live+ema"}},
+        "load": {
+            "default_weights": "ema",
+            "entrypoint": ".",
+            "format": "solarwm_wan_stage2_transaction_v1",
+        },
+    }
+    (tmp_path / "release-manifest.json").write_text(json.dumps(manifest), encoding="utf-8")
+
+    assert _published_default_weight_role(tmp_path) == "ema"
+
+    manifest["load"]["default_weights"] = "unknown"
+    (tmp_path / "release-manifest.json").write_text(json.dumps(manifest), encoding="utf-8")
+    with pytest.raises(BackendContractError, match="default weight role is invalid"):
+        _published_default_weight_role(tmp_path)
+
+
 def test_stage2_unconditional_matches_conditional_dtype() -> None:
     torch = pytest.importorskip("torch")
 
@@ -432,6 +524,7 @@ def to(self, *args: object, **kwargs: object) -> _Movable:
         "_verified_stage2_inference_checkpoint",
         lambda *_: (tmp_path / "model.pt", "b" * 64, 200),
     )
+    monkeypatch.setattr(stage2, "_published_default_weight_role", lambda *_: "ema")
     monkeypatch.setattr(
         stage2.WanAssetLayout,
         "from_config",

tests/backends/wan22/test_vae_streaming.py

@@ -0,0 +1,40 @@
+from __future__ import annotations
+
+import pytest
+
+from solarwm.backends.wan22.runtime.components import Wan5BVAE
+
+
+def test_wan5b_streaming_decode_keeps_one_cache_across_temporal_tiles() -> None:
+    torch = pytest.importorskip("torch")
+
+    class _Module:
+        def __init__(self) -> None:
+            self.clear_calls = 0
+            self.chunk_sizes: list[int] = []
+            self.first = True
+
+        def clear_cache(self) -> None:
+            self.clear_calls += 1
+            self.first = True
+
+        def cached_decode(self, value: object, _scale: object) -> object:
+            latent_frames = int(value.shape[2])
+            self.chunk_sizes.append(latent_frames)
+            pixel_frames = 1 + 4 * (latent_frames - 1) if self.first else 4 * latent_frames
+            self.first = False
+            return torch.zeros((1, 3, pixel_frames, 1, 1), dtype=torch.bfloat16)
+
+    module = _Module()
+    vae = Wan5BVAE.__new__(Wan5BVAE)
+    vae.module = module
+    vae._mean = torch.zeros(48, dtype=torch.float32)
+    vae._std = torch.ones(48, dtype=torch.float32)
+    latents = torch.zeros((1, 240, 48, 1, 1), dtype=torch.bfloat16)
+
+    decoded = vae.decode_streaming(latents, chunk_latent_frames=60)
+
+    assert module.chunk_sizes == [60, 60, 60, 60]
+    assert module.clear_calls == 2
+    assert decoded.device.type == "cpu"
+    assert decoded.shape == (1, 957, 3, 1, 1)