Lazily materialize camera-length video inputs

Author: Junchao-cs

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

@@ -74,7 +74,31 @@ class _MaterializedCandidate:
     noise_seed: int
     start_adjusted: bool
     pass_rollouts: Mapping[str, int]
-    prepared: _PreparedCase
+    source_pixel_frames: int
+    camera_fingerprint: str
+    prepared: _PreparedCase | None
+
+
+@dataclass(frozen=True)
+class _DeferredCameraCase:
+    row: IndexRow
+    sample_plan: SamplePlan
+    noise_seed: int
+    pass_rollouts: Mapping[str, int]
+    source_pixel_frames: int
+    prompt: str
+    camera_fingerprint: str
+
+
+@dataclass
+class _DeferredCameraInputs:
+    rows: tuple[IndexRow, ...]
+    resolver: Any
+    runtime_guards: CameraGuards
+    manifest_guards: CameraGuards
+    cases: dict[int, _DeferredCameraCase]
+    shards: TarShardReader | None = None
+    reader: RawSampleReader | None = None
 
 
 class _FamilyAdapter:
@@ -600,6 +624,7 @@ def __init__(self, config: Mapping[str, Any], plan: GenerationPlan) -> None:
         self.vae.to(self.device)
         self._loaded_role: str | None = None
         self._prepared: dict[int, _PreparedCase] = {}
+        self._deferred_camera_inputs: _DeferredCameraInputs | None = None
 
     def sync(self) -> None:
         import torch.distributed as dist
@@ -610,7 +635,91 @@ def sync(self) -> None:
     def close(self) -> None:
         from .distributed import cleanup_torchrun
 
-        cleanup_torchrun()
+        try:
+            self._close_deferred_camera_inputs()
+        finally:
+            cleanup_torchrun()
+
+    def _close_deferred_camera_inputs(self) -> None:
+        deferred = getattr(self, "_deferred_camera_inputs", None)
+        if deferred is None:
+            return
+        self._deferred_camera_inputs = None
+        try:
+            if deferred.shards is not None:
+                deferred.shards.close()
+        finally:
+            for slot in deferred.cases:
+                self._prepared.pop(slot, None)
+            deferred.cases.clear()
+            deferred.reader = None
+            deferred.shards = None
+
+    def _materialize_deferred_camera_case(self, case: InferenceCase) -> None:
+        deferred = getattr(self, "_deferred_camera_inputs", None)
+        if deferred is None:
+            raise BackendContractError("Wan adapter has no deferred camera inputs")
+        descriptor = deferred.cases.get(case.slot)
+        if descriptor is None:
+            raise BackendContractError(
+                f"Wan adapter has no deferred camera case {case.sample_id!r}"
+            )
+        if self._prepared:
+            raise BackendContractError("Wan adapter retained another prepared camera case")
+        metadata = case.metadata
+        if (
+            descriptor.row.sample_id != case.sample_id
+            or descriptor.sample_plan.start_frame != case.start_frame
+            or descriptor.noise_seed != case.noise_seed
+            or descriptor.prompt != case.prompt
+            or descriptor.camera_fingerprint != case.camera_fingerprint
+            or descriptor.row.key != str(metadata.get("key", ""))
+            or descriptor.row.shard != str(metadata.get("source_shard", ""))
+            or descriptor.row.ordinal != int(metadata.get("source_row_ordinal", -1))
+            or descriptor.source_pixel_frames != int(metadata.get("source_pixel_frames", -1))
+            or descriptor.sample_plan.source_frame_indices[-1]
+            != int(metadata.get("source_frame_last", -1))
+            or dict(descriptor.pass_rollouts)
+            != dict(metadata.get("rollout_latent_frames_by_pass", {}))
+        ):
+            raise DataContractError(
+                f"deferred camera case {case.sample_id!r} identity drifted"
+            )
+        if deferred.reader is None:
+            deferred.shards = TarShardReader(
+                deferred.resolver,
+                max_open=int(self.config["data"].get("tar_cache_size", 4)),
+            )
+            deferred.reader = RawSampleReader(deferred.rows, deferred.shards)
+        raw = deferred.reader.materialize(descriptor.sample_plan)
+        camera_fingerprint = hashlib.blake2s(raw.members["camera_member"]).hexdigest()
+        if raw.caption != case.prompt or camera_fingerprint != case.camera_fingerprint:
+            raise DataContractError(
+                f"deferred camera case {case.sample_id!r} payload drifted"
+            )
+        data = self.config["data"]
+        camera = build_camera_tokens(
+            raw.members["camera_member"],
+            descriptor.sample_plan.source_frame_indices,
+            raw.manifest,
+            source_fps=_source_fps(descriptor.row),
+            output_fps=float(data.get("fps", 16.0)),
+            frame_sequence_length=int(self.config["model"]["frame_sequence_length"]),
+            guards=deferred.runtime_guards,
+            manifest_guards=deferred.manifest_guards,
+            configured_array_key=str(data["camera_array_key"]),
+        )
+        pixels = decode_video(
+            raw.members["video_member"],
+            descriptor.sample_plan.source_frame_indices,
+            height=int(data["height"]),
+            width=int(data["width"]),
+        )
+        self._prepared[case.slot] = _PreparedCase(
+            pixels=pixels,
+            camera=camera,
+            source_pixel_frames=descriptor.source_pixel_frames,
+        )
 
     def weight_id(self, role: str) -> str:
         if role not in {"live", "ema"}:
@@ -669,6 +778,11 @@ def _load_role(self, role: str) -> None:
     def build_cases(self, plan: GenerationPlan) -> tuple[InferenceCase, ...]:
         data = self.config["data"]
         validation = self.config["validation"]
+        camera_length = (
+            str(self.config.get("inference", {}).get("length", "fixed")).strip().lower() == "camera"
+        )
+        if camera_length:
+            self._close_deferred_camera_inputs()
         frozen_path_value = getattr(self, "validation_plan_path", None)
         frozen_path = Path(str(frozen_path_value)) if frozen_path_value is not None else None
         frozen_key = validation_plan_key("wan22", self.config)
@@ -724,9 +838,14 @@ def validation_guard(name: str) -> float | None:
             max_rel_translation=float(data["max_rel_translation"]),
             max_camera_abs=float(data["max_camera_abs"]),
         )
-        camera_length = (
-            str(self.config.get("inference", {}).get("length", "fixed")).strip().lower() == "camera"
-        )
+        if camera_length:
+            self._deferred_camera_inputs = _DeferredCameraInputs(
+                rows=tuple(source_rows),
+                resolver=resolver,
+                runtime_guards=runtime_guards,
+                manifest_guards=manifest_guards,
+                cases={},
+            )
         variable_rollout = camera_length or any(
             item.variable_rollout_by_source for item in plan.passes
         )
@@ -742,6 +861,11 @@ def validation_guard(name: str) -> float | None:
             # Candidate rows retain full-index ordinals, so random access must
             # address the complete source index rather than the seeded order.
             reader = RawSampleReader(source_rows, shards)
+            camera_reader = RawSampleReader(
+                source_rows,
+                shards,
+                member_fields=("camera_member",),
+            )
 
             def materialize_candidate(row: IndexRow, slot: int) -> _MaterializedCandidate:
                 pass_rollouts = {
@@ -765,13 +889,15 @@ def materialize_candidate(row: IndexRow, slot: int) -> _MaterializedCandidate:
                     variable_rollout_by_source=variable_rollout,
                     start_at_first_frame=camera_length,
                 )
-                raw = reader.materialize(sample_plan)
-                pixels = decode_video(
-                    raw.members["video_member"],
-                    sample_plan.source_frame_indices,
-                    height=int(data["height"]),
-                    width=int(data["width"]),
-                )
+                raw = (camera_reader if camera_length else reader).materialize(sample_plan)
+                pixels = None
+                if not camera_length:
+                    pixels = decode_video(
+                        raw.members["video_member"],
+                        sample_plan.source_frame_indices,
+                        height=int(data["height"]),
+                        width=int(data["width"]),
+                    )
                 camera = build_camera_tokens(
                     raw.members["camera_member"],
                     sample_plan.source_frame_indices,
@@ -789,13 +915,44 @@ def materialize_candidate(row: IndexRow, slot: int) -> _MaterializedCandidate:
                     noise_seed=noise_seed,
                     start_adjusted=start_adjusted,
                     pass_rollouts=pass_rollouts,
-                    prepared=_PreparedCase(
-                        pixels=pixels,
-                        camera=camera,
-                        source_pixel_frames=source_pixel_frames,
+                    source_pixel_frames=source_pixel_frames,
+                    camera_fingerprint=hashlib.blake2s(
+                        raw.members["camera_member"]
+                    ).hexdigest(),
+                    prepared=(
+                        None
+                        if camera_length
+                        else _PreparedCase(
+                            pixels=pixels,
+                            camera=camera,
+                            source_pixel_frames=source_pixel_frames,
+                        )
                     ),
                 )
 
+            def retain_candidate(
+                row: IndexRow,
+                slot: int,
+                materialized: _MaterializedCandidate,
+            ) -> None:
+                if camera_length:
+                    deferred = self._deferred_camera_inputs
+                    if deferred is None or materialized.prepared is not None:
+                        raise BackendContractError("Wan deferred camera state is invalid")
+                    deferred.cases[slot] = _DeferredCameraCase(
+                        row=row,
+                        sample_plan=materialized.sample_plan,
+                        noise_seed=materialized.noise_seed,
+                        pass_rollouts=dict(materialized.pass_rollouts),
+                        source_pixel_frames=materialized.source_pixel_frames,
+                        prompt=str(materialized.raw.caption),
+                        camera_fingerprint=materialized.camera_fingerprint,
+                    )
+                    return
+                if materialized.prepared is None:
+                    raise BackendContractError("Wan fixed validation case lacks pixels")
+                self._prepared[slot] = materialized.prepared
+
             if frozen_cases is not None:
                 restored: list[InferenceCase] = []
                 for case in frozen_cases:
@@ -821,21 +978,18 @@ def materialize_candidate(row: IndexRow, slot: int) -> _MaterializedCandidate:
                         raise DataContractError(
                             f"frozen validation sample {case.sample_id!r} changed camera validity"
                         ) from exc
-                    camera_fingerprint = hashlib.blake2s(
-                        materialized.raw.members["camera_member"]
-                    ).hexdigest()
                     if (
                         materialized.sample_plan.start_frame != case.start_frame
                         or materialized.noise_seed != case.noise_seed
                         or materialized.raw.caption != case.prompt
-                        or camera_fingerprint != case.camera_fingerprint
+                        or materialized.camera_fingerprint != case.camera_fingerprint
                         or dict(materialized.pass_rollouts)
                         != dict(case.metadata.get("rollout_latent_frames_by_pass", {}))
                     ):
                         raise DataContractError(
                             f"frozen validation sample {case.sample_id!r} materialization drifted"
                         )
-                    self._prepared[case.slot] = materialized.prepared
+                    retain_candidate(row, case.slot, materialized)
                     restored.append(case)
                 return tuple(restored)
 
@@ -930,17 +1084,15 @@ def materialize_candidate(row: IndexRow, slot: int) -> _MaterializedCandidate:
                         prompt=raw.caption,
                         start_frame=sample_plan.start_frame,
                         noise_seed=materialized.noise_seed,
-                        camera_fingerprint=hashlib.blake2s(
-                            raw.members["camera_member"]
-                        ).hexdigest(),
+                        camera_fingerprint=materialized.camera_fingerprint,
                         metadata={
                             "key": row.key,
                             "source_shard": row.shard,
                             "source_row_ordinal": row.ordinal,
                             "dataset": _row_dataset(row),
                             "scene": row.values.get("scene"),
                             "source_num_frames": _num_frames(row),
-                            "source_pixel_frames": materialized.prepared.source_pixel_frames,
+                            "source_pixel_frames": materialized.source_pixel_frames,
                             "train_latent_frames": int(data["latent_frames"]),
                             "rollout_latent_frames_by_pass": materialized.pass_rollouts,
                             **({"rollout_length_source": "camera"} if camera_length else {}),
@@ -955,7 +1107,7 @@ def materialize_candidate(row: IndexRow, slot: int) -> _MaterializedCandidate:
                         },
                     )
                 )
-                self._prepared[slot] = materialized.prepared
+                retain_candidate(row, slot, materialized)
         local_cases = tuple(cases)
         if frozen_path is not None:
             all_cases = _gather_partitioned_cases(local_cases)

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

@@ -2298,7 +2298,8 @@ def __init__(self, values: Mapping[str, Any], generation_plan: Any) -> None:
                 self.text_encoder.to(self.device)
                 self.vae.to(self.device)
                 self._loaded_role: str | None = None
-                self._prepared: dict[str, Any] = {}
+                self._prepared: dict[int, Any] = {}
+                self._deferred_camera_inputs = None
 
             @staticmethod
             def _root(module: Any) -> Any:
@@ -2343,7 +2344,22 @@ def generate(self, case: Any, *, weights_id: str) -> Any:
                         "Stage2 standalone generation weights identity drift"
                     )
                 self._load_role(role)
-                return _stage2_generated_sample(self, case, weights_id=weights_id)
+                camera_length = (
+                    str(self.config.get("inference", {}).get("length", "fixed"))
+                    .strip()
+                    .lower()
+                    == "camera"
+                )
+                deferred = self._deferred_camera_inputs
+                if camera_length and deferred is None:
+                    raise BackendContractError("Stage2 camera-length inputs were not prepared")
+                if not camera_length:
+                    return _stage2_generated_sample(self, case, weights_id=weights_id)
+                self._materialize_deferred_camera_case(case)
+                try:
+                    return _stage2_generated_sample(self, case, weights_id=weights_id)
+                finally:
+                    self._prepared.pop(case.slot, None)
 
         return _Adapter(config, plan)