anyflow

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</code></pre></details>