pipeline/predictor_v4_dmd_training.py

"""Random-exit DMD rollout for Predictor-v4 and optional Full-Generator tuning."""

from __future__ import annotations

from contextlib import nullcontext
from typing import Any

import torch
import torch.distributed as dist

from predictor_training.rollout_cache import (
    assert_clean_history_extent,
    build_predictor_workspace,
    reset_main_caches,
    reset_predictor_workspace,
)


CHUNK_FRAMES = 3
NUM_CHUNKS = 7
NUM_STEPS = 4
TOKENS_PER_FRAME = 1560
TOKENS_PER_CHUNK = CHUNK_FRAMES * TOKENS_PER_FRAME


def _unwrap(module: torch.nn.Module) -> torch.nn.Module:
    while hasattr(module, "module"):
        module = module.module
    return module


class PredictorV4DMDTrainingPipeline:
    """Generate one detached random-exit trajectory for DMD.

    ``predictor_only`` samples exits P1/P2/P3. ``joint`` additionally samples
    the first Full step.  One exit is shared by all chunks and distributed
    ranks.  Only the selected call for each chunk retains a graph.
    """

    def __init__(
        self,
        *,
        denoising_step_list: torch.Tensor,
        scheduler: Any,
        generator: torch.nn.Module,
        predictor: torch.nn.Module,
        training_mode: str,
        context_noise: int = 0,
        forced_exit_step: int | None = None,
    ) -> None:
        training_mode = str(training_mode).lower()
        if training_mode not in {"predictor_only", "joint"}:
            raise ValueError(f"Unknown Predictor DMD mode {training_mode!r}")
        if len(denoising_step_list) != NUM_STEPS:
            raise ValueError("Predictor-v4 DMD requires exactly four denoising steps")
        self.denoising_step_list = denoising_step_list
        self.scheduler = scheduler
        self.generator = generator
        self.predictor = predictor
        self.training_mode = training_mode
        self.context_noise = int(context_noise)
        self.forced_exit_step = (
            None if forced_exit_step is None else int(forced_exit_step)
        )
        if self.forced_exit_step is not None:
            minimum = 1 if training_mode == "predictor_only" else 0
            if not minimum <= self.forced_exit_step < NUM_STEPS:
                raise ValueError(
                    f"forced_exit_step must be in [{minimum}, {NUM_STEPS - 1}]"
                )
        self.kv_cache1: list[dict[str, torch.Tensor]] | None = None
        self.crossattn_cache: list[dict[str, torch.Tensor | bool]] | None = None
        self.last_exit_step: int | None = None

    @property
    def predictor_module(self) -> torch.nn.Module:
        return _unwrap(self.predictor)

    @property
    def generator_wrapper(self) -> torch.nn.Module:
        return _unwrap(self.generator)

    def _sample_exit_step(self, device: torch.device) -> int:
        if self.forced_exit_step is not None:
            value = torch.tensor(
                [self.forced_exit_step], device=device, dtype=torch.long
            )
        elif not dist.is_initialized() or dist.get_rank() == 0:
            low = 1 if self.training_mode == "predictor_only" else 0
            value = torch.randint(low, NUM_STEPS, (1,), device=device)
        else:
            value = torch.empty(1, device=device, dtype=torch.long)
        if dist.is_initialized():
            dist.broadcast(value, src=0)
        return int(value.item())

    def _initialize_caches(
        self,
        *,
        batch_size: int,
        dtype: torch.dtype,
        device: torch.device,
    ) -> None:
        capacity = NUM_CHUNKS * TOKENS_PER_CHUNK
        self.kv_cache1 = [
            {
                "k": torch.zeros(
                    batch_size, capacity, 12, 128, dtype=dtype, device=device
                ),
                "v": torch.zeros(
                    batch_size, capacity, 12, 128, dtype=dtype, device=device
                ),
                "global_end_index": torch.zeros(
                    1, dtype=torch.long, device=device
                ),
                "local_end_index": torch.zeros(
                    1, dtype=torch.long, device=device
                ),
            }
            for _ in range(30)
        ]
        self.crossattn_cache = [
            {
                "k": torch.zeros(
                    batch_size, 512, 12, 128, dtype=dtype, device=device
                ),
                "v": torch.zeros(
                    batch_size, 512, 12, 128, dtype=dtype, device=device
                ),
                "is_init": False,
            }
            for _ in range(30)
        ]

    def _prepare_caches(self, noise: torch.Tensor) -> None:
        if (
            self.kv_cache1 is None
            or self.kv_cache1[0]["k"].shape[0] != noise.shape[0]
            or self.kv_cache1[0]["k"].device != noise.device
            or self.kv_cache1[0]["k"].dtype != noise.dtype
        ):
            self._initialize_caches(
                batch_size=noise.shape[0],
                dtype=noise.dtype,
                device=noise.device,
            )
        else:
            reset_main_caches(self.kv_cache1, self.crossattn_cache)

    def _timestep(
        self,
        step_id: int,
        *,
        batch_size: int,
        device: torch.device,
    ) -> torch.Tensor:
        return torch.full(
            (batch_size, CHUNK_FRAMES),
            float(self.denoising_step_list[int(step_id)].item()),
            dtype=torch.float32,
            device=device,
        )

    def _renoise(
        self,
        clean: torch.Tensor,
        *,
        next_timestep: torch.Tensor,
    ) -> torch.Tensor:
        return self.scheduler.add_noise(
            clean.detach().flatten(0, 1),
            torch.randn_like(clean.detach().flatten(0, 1)),
            next_timestep.flatten(0, 1),
        ).unflatten(0, clean.shape[:2]).detach()

    def _full_step(
        self,
        *,
        latent: torch.Tensor,
        conditional_dict: dict[str, torch.Tensor],
        timestep: torch.Tensor,
        current_start: int,
        capture_hidden: bool,
        retain_graph: bool,
    ) -> tuple[torch.Tensor, torch.Tensor, torch.Tensor | None]:
        captured: list[torch.Tensor] = []
        handle = None
        if capture_hidden:
            def capture_head_input(module, args):
                if not args or not torch.is_tensor(args[0]):
                    raise RuntimeError("Wan head hook did not receive hidden states")
                captured.append(args[0])

            handle = self.generator_wrapper.model.head.register_forward_pre_hook(
                capture_head_input
            )
        grad_context = nullcontext() if retain_graph else torch.no_grad()
        try:
            with grad_context, torch.autocast(
                device_type="cuda", dtype=torch.bfloat16
            ):
                flow, clean = self.generator(
                    noisy_image_or_video=latent,
                    conditional_dict=conditional_dict,
                    timestep=timestep,
                    kv_cache=self.kv_cache1,
                    crossattn_cache=self.crossattn_cache,
                    current_start=current_start,
                )
        finally:
            if handle is not None:
                handle.remove()
        if capture_hidden and len(captured) != 1:
            raise RuntimeError(
                f"Expected one Full hidden capture, got {len(captured)}"
            )
        return flow, clean, captured[0] if captured else None

    def _predictor_step(
        self,
        *,
        latent: torch.Tensor,
        timestep: torch.Tensor,
        anchor_hidden: torch.Tensor,
        previous_chunk_hidden: torch.Tensor,
        predictor_cache: dict[int, dict[str, torch.Tensor]],
        current_start: int,
        retain_graph: bool,
    ) -> dict[str, torch.Tensor]:
        grad_context = nullcontext() if retain_graph else torch.no_grad()
        with grad_context, torch.autocast(
            device_type="cuda", dtype=torch.bfloat16
        ):
            return self.predictor(
                target_latent=latent,
                target_timestep=timestep,
                anchor_hidden=anchor_hidden,
                previous_chunk_hidden=previous_chunk_hidden,
                kv_cache=predictor_cache,
                crossattn_cache=self.crossattn_cache,
                current_start=current_start,
            )

    @torch.no_grad()
    def _commit_clean_history(
        self,
        *,
        clean: torch.Tensor,
        conditional_dict: dict[str, torch.Tensor],
        current_start: int,
        expected_end: int,
    ) -> None:
        timestep = torch.full(
            (clean.shape[0], clean.shape[1]),
            float(self.context_noise),
            dtype=torch.float32,
            device=clean.device,
        )
        with torch.autocast(device_type="cuda", dtype=torch.bfloat16):
            self.generator(
                noisy_image_or_video=clean.detach(),
                conditional_dict=conditional_dict,
                timestep=timestep,
                kv_cache=self.kv_cache1,
                crossattn_cache=self.crossattn_cache,
                current_start=current_start,
            )
        assert_clean_history_extent(
            self.kv_cache1, expected_tokens=expected_end
        )

    def inference_with_trajectory(
        self,
        noise: torch.Tensor,
        initial_latent: torch.Tensor | None = None,
        return_sim_step: bool = False,
        **conditional_dict: torch.Tensor,
    ):
        if initial_latent is not None:
            raise NotImplementedError("Predictor-v4 DMD currently supports T2V only")
        if noise.shape[1] != NUM_CHUNKS * CHUNK_FRAMES:
            raise ValueError(
                f"Expected 21 latent frames, got shape {tuple(noise.shape)}"
            )
        self._prepare_caches(noise)
        exit_step = self._sample_exit_step(noise.device)
        self.last_exit_step = exit_step
        outputs: list[torch.Tensor] = []
        previous_hidden: dict[int, torch.Tensor] = {}

        for chunk_id in range(NUM_CHUNKS):
            current_start = chunk_id * TOKENS_PER_CHUNK
            current_end = current_start + TOKENS_PER_CHUNK
            latent = noise[
                :, chunk_id * CHUNK_FRAMES : (chunk_id + 1) * CHUNK_FRAMES
            ].detach()

            if exit_step == 0:
                timestep = self._timestep(
                    0, batch_size=noise.shape[0], device=noise.device
                )
                _, clean, _ = self._full_step(
                    latent=latent,
                    conditional_dict=conditional_dict,
                    timestep=timestep,
                    current_start=current_start,
                    capture_hidden=False,
                    retain_graph=self.training_mode == "joint",
                )
                outputs.append(clean)
                self._commit_clean_history(
                    clean=clean,
                    conditional_dict=conditional_dict,
                    current_start=current_start,
                    expected_end=current_end,
                )
                continue

            if chunk_id == 0:
                current_hidden: dict[int, torch.Tensor] = {}
                clean = None
                for step_id in range(exit_step + 1):
                    timestep = self._timestep(
                        step_id,
                        batch_size=noise.shape[0],
                        device=noise.device,
                    )
                    _, clean, hidden = self._full_step(
                        latent=latent,
                        conditional_dict=conditional_dict,
                        timestep=timestep,
                        current_start=current_start,
                        capture_hidden=step_id > 0,
                        retain_graph=(
                            self.training_mode == "joint"
                            and step_id == exit_step
                        ),
                    )
                    if step_id > 0:
                        if hidden is None:
                            raise RuntimeError(
                                f"Chunk 0 Full step {step_id} has no hidden"
                            )
                        current_hidden[step_id] = hidden.detach()
                    if step_id < exit_step:
                        latent = self._renoise(
                            clean,
                            next_timestep=self._timestep(
                                step_id + 1,
                                batch_size=noise.shape[0],
                                device=noise.device,
                            ),
                        )
                if clean is None:
                    raise RuntimeError("Chunk 0 produced no clean latent")
                outputs.append(clean)
                self._commit_clean_history(
                    clean=clean,
                    conditional_dict=conditional_dict,
                    current_start=current_start,
                    expected_end=current_end,
                )
                previous_hidden = current_hidden
                continue

            predictor_cache = build_predictor_workspace(
                self.kv_cache1,
                source_block_ids=tuple(self.predictor_module.source_block_ids),
                history_tokens=current_start,
                current_tokens=TOKENS_PER_CHUNK,
            )
            timestep0 = self._timestep(
                0, batch_size=noise.shape[0], device=noise.device
            )
            _, clean0, anchor = self._full_step(
                latent=latent,
                conditional_dict=conditional_dict,
                timestep=timestep0,
                current_start=current_start,
                capture_hidden=True,
                retain_graph=False,
            )
            if anchor is None:
                raise RuntimeError(f"Chunk {chunk_id} has no Full step-0 anchor")
            latent = self._renoise(
                clean0,
                next_timestep=self._timestep(
                    1, batch_size=noise.shape[0], device=noise.device
                ),
            )
            current_hidden = {}
            pred_clean = None

            for step_id in range(1, exit_step + 1):
                if step_id not in previous_hidden:
                    raise RuntimeError(
                        f"Previous chunk lacks Predictor hidden step {step_id}"
                    )
                reset_predictor_workspace(
                    predictor_cache, history_tokens=current_start
                )
                timestep = self._timestep(
                    step_id,
                    batch_size=noise.shape[0],
                    device=noise.device,
                )
                output = self._predictor_step(
                    latent=latent,
                    timestep=timestep,
                    anchor_hidden=anchor.detach(),
                    previous_chunk_hidden=previous_hidden[step_id],
                    predictor_cache=predictor_cache,
                    current_start=current_start,
                    retain_graph=step_id == exit_step,
                )
                flow = output["pred_flow"]
                pred_clean = self.generator_wrapper._convert_flow_pred_to_x0(
                    flow_pred=flow.flatten(0, 1),
                    xt=latent.flatten(0, 1),
                    timestep=timestep.flatten(0, 1),
                ).unflatten(0, flow.shape[:2])
                anchor = output["pred_hidden"].detach()
                current_hidden[step_id] = anchor
                if step_id < exit_step:
                    latent = self._renoise(
                        pred_clean,
                        next_timestep=self._timestep(
                            step_id + 1,
                            batch_size=noise.shape[0],
                            device=noise.device,
                        ),
                    )

            if pred_clean is None:
                raise RuntimeError(f"Chunk {chunk_id} produced no Predictor output")
            outputs.append(pred_clean)
            self._commit_clean_history(
                clean=pred_clean,
                conditional_dict=conditional_dict,
                current_start=current_start,
                expected_end=current_end,
            )
            previous_hidden = current_hidden

        result = torch.cat(outputs, dim=1)
        if return_sim_step:
            return result, 0, 0, exit_step + 1
        return result, 0, 0


__all__ = ["PredictorV4DMDTrainingPipeline"]