210 lines
8.7 KiB
Python
210 lines
8.7 KiB
Python
# SPDX-License-Identifier: Apache-2.0
|
|
"""AnyFlow on-policy distillation method.
|
|
|
|
Stage 2 of the AnyFlow two-stage recipe. Continues from a pretrained
|
|
flow-map student and refines it via distribution-matching distillation
|
|
(DMD2) where the student is rolled out for ``student_sample_steps``
|
|
Euler-flow steps from pure noise. One randomly-chosen step in the
|
|
rollout is gradient-enabled (and broadcast across ranks so every worker
|
|
agrees on which step to gradient-enable); the rest run under
|
|
``torch.no_grad``.
|
|
|
|
Inherits ``DMD2Method`` for the alternating student / critic update
|
|
machinery and the existing DMD VSD-with-fake-score loss. Overrides
|
|
``_student_rollout`` to drive the multi-step Euler-flow rollout with
|
|
``r = t_next`` (mean-velocity sampling — matches the AnyFlow paper's
|
|
``WanAnyFlowPipeline.training_rollout`` with ``use_mean_velocity=True``).
|
|
|
|
Reference: ``pipeline_wan_anyflow.py::training_rollout`` in
|
|
NVlabs/AnyFlow at commit ``549236a``.
|
|
"""
|
|
|
|
from __future__ import annotations
|
|
|
|
from typing import Any, Literal
|
|
|
|
import torch
|
|
import torch.distributed as dist
|
|
|
|
from fastvideo.train.methods.distribution_matching.dmd2 import DMD2Method
|
|
from fastvideo.train.utils.config import (
|
|
get_optional_float,
|
|
get_optional_int,
|
|
)
|
|
|
|
|
|
class AnyFlowMethod(DMD2Method):
|
|
"""AnyFlow on-policy distillation (multi-step rollout)."""
|
|
|
|
def __init__(
|
|
self,
|
|
*,
|
|
cfg: Any,
|
|
role_models: dict[str, Any],
|
|
) -> None:
|
|
super().__init__(cfg=cfg, role_models=role_models)
|
|
mcfg = self.method_config
|
|
|
|
student_sample_steps = get_optional_int(mcfg, "student_sample_steps", where="method.student_sample_steps")
|
|
if student_sample_steps is None:
|
|
student_sample_steps = 4
|
|
if int(student_sample_steps) <= 0:
|
|
raise ValueError("method.student_sample_steps must be positive, "
|
|
f"got {student_sample_steps}")
|
|
self._student_sample_steps = int(student_sample_steps)
|
|
|
|
use_mean_velocity_raw = mcfg.get("use_mean_velocity", True)
|
|
if not isinstance(use_mean_velocity_raw, bool):
|
|
raise ValueError("method.use_mean_velocity must be a bool, "
|
|
f"got {type(use_mean_velocity_raw).__name__}")
|
|
self._use_mean_velocity = bool(use_mean_velocity_raw)
|
|
|
|
# Optional pinned rollout schedule (descending, absolute t-units).
|
|
# Falls back to dmd_denoising_steps when absent.
|
|
raw_t_list = mcfg.get("t_list_override", None)
|
|
if raw_t_list is None:
|
|
self._t_list_override: list[float] | None = None
|
|
else:
|
|
if not isinstance(raw_t_list, list) or not raw_t_list:
|
|
raise ValueError("method.t_list_override must be a non-empty list of "
|
|
f"floats when set, got {raw_t_list!r}")
|
|
t_list = [float(x) for x in raw_t_list]
|
|
for i in range(len(t_list) - 1):
|
|
if t_list[i] < t_list[i + 1]:
|
|
raise ValueError("method.t_list_override must be descending, "
|
|
f"got {t_list!r}")
|
|
self._t_list_override = t_list
|
|
|
|
# Scoring conditioning: AnyFlow scores against r=0 for the DMD branch.
|
|
score_r_raw = mcfg.get("dmd_score_r_value", 0.0)
|
|
try:
|
|
self._dmd_score_r = float(score_r_raw)
|
|
except (TypeError, ValueError) as exc:
|
|
raise ValueError("method.dmd_score_r_value must be numeric, "
|
|
f"got {score_r_raw!r}") from exc
|
|
|
|
# Optional teacher guidance scale for the DMD loss (carry over from
|
|
# DMD2Method's behavior; default 1.0).
|
|
guidance = get_optional_float(mcfg, "real_score_guidance_scale", where="method.real_score_guidance_scale")
|
|
self._real_score_guidance = float(guidance) if guidance is not None else 1.0
|
|
|
|
# ------------------------------------------------------------------
|
|
# Rollout schedule
|
|
|
|
def _get_rollout_schedule(self, *, device: torch.device) -> torch.Tensor:
|
|
"""Build the descending timestep schedule used by the on-policy
|
|
rollout. Length is ``num_steps + 1`` so ``num_steps`` Euler steps
|
|
consume the full range.
|
|
|
|
Order of precedence:
|
|
1. ``method.t_list_override`` — used verbatim (absolute units).
|
|
2. ``method.dmd_denoising_steps`` (inherited from DMD2) appended
|
|
with a final 0 boundary if the last entry isn't already 0.
|
|
"""
|
|
if self._t_list_override is not None:
|
|
return torch.tensor(self._t_list_override, device=device, dtype=torch.float32)
|
|
|
|
steps = self._get_denoising_step_list(device).to(dtype=torch.float32)
|
|
if float(steps[-1].item()) != 0.0:
|
|
zero = torch.zeros(1, device=device, dtype=torch.float32)
|
|
steps = torch.cat([steps, zero], dim=0)
|
|
return steps
|
|
|
|
def _broadcast_grad_step_index(
|
|
self,
|
|
num_steps: int,
|
|
*,
|
|
device: torch.device,
|
|
) -> int:
|
|
"""Pick the rollout step that gets gradient enabled. In distributed
|
|
runs the choice is broadcast from rank 0 so every worker agrees."""
|
|
if num_steps <= 0:
|
|
raise ValueError("num_steps must be positive")
|
|
if dist.is_initialized() and dist.get_rank() != 0:
|
|
idx_tensor = torch.empty((1, ), dtype=torch.long, device=device)
|
|
else:
|
|
idx_tensor = torch.randint(0,
|
|
num_steps, (1, ),
|
|
device=device,
|
|
dtype=torch.long,
|
|
generator=self.cuda_generator)
|
|
if dist.is_initialized():
|
|
dist.broadcast(idx_tensor, src=0)
|
|
return int(idx_tensor.item())
|
|
|
|
# ------------------------------------------------------------------
|
|
# Rollout
|
|
|
|
def _student_rollout(
|
|
self,
|
|
batch: Any,
|
|
*,
|
|
with_grad: bool,
|
|
) -> torch.Tensor:
|
|
"""Multi-step Euler-flow rollout from pure noise.
|
|
|
|
Returns the predicted clean latent ``x_0`` after the chosen
|
|
gradient step (or the final ``x`` after the last step if
|
|
``with_grad`` is False — used by the critic path).
|
|
"""
|
|
latents = batch.latents
|
|
if latents is None or latents.ndim != 5:
|
|
raise RuntimeError("AnyFlow on-policy rollout requires TrainingBatch.latents "
|
|
"of shape [B, T, C, H, W] for shape templating")
|
|
device = latents.device
|
|
dtype = latents.dtype
|
|
|
|
schedule = self._get_rollout_schedule(device=device)
|
|
num_entries = int(schedule.numel())
|
|
num_steps = num_entries - 1
|
|
if num_steps <= 0:
|
|
raise RuntimeError("rollout schedule must have at least two entries "
|
|
f"(got {num_entries})")
|
|
if num_steps > self._student_sample_steps:
|
|
# Trim to the configured cap, keeping the last (=0) boundary.
|
|
schedule = torch.cat([schedule[:self._student_sample_steps], schedule[-1:]], dim=0)
|
|
num_steps = self._student_sample_steps
|
|
|
|
grad_step = self._broadcast_grad_step_index(num_steps, device=device) if with_grad else -1
|
|
|
|
attn_kind: Literal["dense", "vsa"] = "vsa"
|
|
n_train = float(self.student.num_train_timesteps)
|
|
|
|
x = torch.randn(latents.shape, device=device, dtype=dtype, generator=self.cuda_generator)
|
|
last_pred_x0: torch.Tensor | None = None
|
|
batch_size = int(latents.shape[0])
|
|
|
|
for i in range(num_steps):
|
|
t_cur = schedule[i].expand(batch_size)
|
|
t_next = schedule[i + 1].expand(batch_size)
|
|
r = t_next if self._use_mean_velocity else t_cur
|
|
|
|
enable_grad = bool(with_grad) and (i == grad_step)
|
|
with torch.set_grad_enabled(enable_grad):
|
|
v = self.student.predict_velocity_with_r(
|
|
x,
|
|
t_cur,
|
|
r,
|
|
batch,
|
|
conditional=True,
|
|
cfg_uncond=self._cfg_uncond,
|
|
attn_kind=attn_kind,
|
|
)
|
|
# Euler step in absolute units: x ← x - ((t_cur - t_next) / N) * v.
|
|
view = [-1] + [1] * (x.ndim - 1)
|
|
dt = ((t_cur - t_next) / n_train).view(*view)
|
|
x = x - dt * v
|
|
|
|
if enable_grad:
|
|
# We treat the rollout output (post-step) as a predicted
|
|
# clean latent — AnyFlow's last Euler step lands at t=0.
|
|
last_pred_x0 = x
|
|
|
|
if last_pred_x0 is None:
|
|
# No gradient step taken (with_grad=False path).
|
|
last_pred_x0 = x
|
|
|
|
if hasattr(batch, "dmd_latent_vis_dict"):
|
|
batch.dmd_latent_vis_dict["generator_timestep"] = (schedule[-1].detach().clone())
|
|
return last_pred_x0
|