anyflow_pretrain

class AnyFlowPretrainMethod(TrainingMethod):
    """AnyFlow flow-map pretrain method.
Single-student training; no teacher or critic. The student must
implement ``predict_velocity_with_r(noisy, t, r, batch, ...)`` —
typically a ``WanModel`` with ``r_embedder=True`` in its arch config.
"""

def __init__(
    self,
    *,
    cfg: Any,
    role_models: dict[str, ModelBase],
) -> None:
    super().__init__(cfg=cfg, role_models=role_models)
    if "student" not in role_models:
        raise ValueError("AnyFlowPretrainMethod requires role 'student'")
    if not self.student._trainable:
        raise ValueError("AnyFlowPretrainMethod requires student to be trainable")

    mcfg = self.method_config
    self._diffusion_ratio = float(
        get_optional_float(mcfg, "diffusion_ratio", where="method.diffusion_ratio") or 0.5)
    self._consistency_ratio = float(
        get_optional_float(mcfg, "consistency_ratio", where="method.consistency_ratio") or 0.25)
    if self._diffusion_ratio + self._consistency_ratio > 1.0:
        raise ValueError("method.diffusion_ratio + method.consistency_ratio must "
                         f"be <= 1, got {self._diffusion_ratio} + "
                         f"{self._consistency_ratio}")

    # δ: finite-difference step in absolute train-timestep units.
    epsilon = get_optional_int(mcfg, "epsilon", where="method.epsilon")
    self._fd_epsilon = float(epsilon) if epsilon is not None else 5.0

    # Loss weighting scheme (uniform / gaussian / beta08).
    raw_weight_type = mcfg.get("weight_type", "beta08")
    if not isinstance(raw_weight_type, str):
        raise ValueError("method.weight_type must be a string, got "
                         f"{type(raw_weight_type).__name__}")
    weight_type = raw_weight_type.strip().lower()
    if weight_type not in {"uniform", "gaussian", "beta08"}:
        raise ValueError("method.weight_type must be one of "
                         "{uniform, gaussian, beta08}, "
                         f"got {raw_weight_type!r}")
    self._weight_type = weight_type

    # Guidance fused into the training target (default 1.0 = unused).
    fg = get_optional_float(mcfg, "fuse_guidance_scale", where="method.fuse_guidance_scale")
    self._fuse_guidance_scale = float(fg) if fg is not None else 1.0
    if self._fuse_guidance_scale <= 0.0:
        raise ValueError("method.fuse_guidance_scale must be positive, "
                         f"got {self._fuse_guidance_scale}")

    # Flow-map scheduler — uses pipeline_config.flow_shift if present
    # and falls back to method.shift (and finally 1.0).
    shift = float(getattr(self.training_config.pipeline_config, "flow_shift", 0.0) or 0.0)
    if shift <= 0.0:
        shift_override = get_optional_float(mcfg, "shift", where="method.shift")
        shift = float(shift_override) if shift_override is not None else 1.0
    self._shift = shift

    # Lazy-imported to avoid circular imports on package load.
    from fastvideo.models.schedulers.scheduling_flow_map_euler_discrete import (
        FlowMapEulerDiscreteScheduler, )
    self._flow_map_scheduler = FlowMapEulerDiscreteScheduler(
        num_train_timesteps=int(self.student.num_train_timesteps),
        shift=self._shift,
    )

    self.student.init_preprocessors(self.training_config)
    self._init_optimizer_and_scheduler()

@property
def _optimizer_dict(self) -> dict[str, torch.optim.Optimizer]:
    return {"student": self._student_optimizer}

@property
def _lr_scheduler_dict(self) -> dict[str, Any]:
    return {"student": self._student_lr_scheduler}

def get_optimizers(
    self,
    iteration: int,
) -> Sequence[torch.optim.Optimizer]:
    del iteration
    return [self._student_optimizer]

def get_lr_schedulers(self, iteration: int) -> Sequence[Any]:
    del iteration
    return [self._student_lr_scheduler]

def single_train_step(
    self,
    batch: dict[str, Any],
    iteration: int,
) -> tuple[
        dict[str, torch.Tensor],
        dict[str, Any],
        dict[str, LogScalar],
]:
    del iteration  # AnyFlow pretrain has no iteration-dependent dispatch.

    training_batch = self.student.prepare_batch(
        batch,
        generator=self.cuda_generator,
        latents_source="data",
    )
    latents = training_batch.latents  # [B, T, C, H, W] (post-permute in prepare_batch).
    if latents is None or latents.ndim != 5:
        raise RuntimeError("AnyFlow pretrain expects TrainingBatch.latents of shape "
                           "[B, T, C, H, W] after prepare_batch; got "
                           f"{None if latents is None else tuple(latents.shape)}")
    device = latents.device
    dtype = latents.dtype
    batch_size = int(latents.shape[0])

    # AnyFlow (t, r) sampling — overrides the timestep drawn by
    # WanModel._sample_timesteps inside prepare_batch.
    t_norm, r_norm, is_diffusion, is_consistency = _sample_pair_timesteps(
        batch_size=batch_size,
        diffusion_ratio=self._diffusion_ratio,
        consistency_ratio=self._consistency_ratio,
        device=device,
        generator=self.cuda_generator,
    )

    sched = self._flow_map_scheduler
    n_train = float(self.student.num_train_timesteps)
    t = (sched.apply_shift(t_norm) * n_train).to(device=device, dtype=dtype)
    r = (sched.apply_shift(r_norm) * n_train).to(device=device, dtype=dtype)

    # Fresh noise drawn from the method's RNG; ignore the noise that
    # prepare_batch attached (it pairs with the discarded timestep).
    noise = torch.randn(
        latents.shape,
        device=device,
        dtype=dtype,
        generator=self.cuda_generator,
    )
    noisy = sched.add_noise(latents, noise, t)

    # Keep training_batch coherent with the new (t, noisy): downstream
    # forward_context uses these fields.
    training_batch.timesteps = t
    training_batch.noise = noise
    training_batch.noisy_model_input = noisy.permute(0, 2, 1, 3, 4)

    # Student velocity prediction at (t, r).
    noise_pred = self.student.predict_velocity_with_r(
        noisy,
        t,
        r,
        training_batch,
        conditional=True,
        attn_kind="dense",
    )

    # Optional guidance distillation — fuse CFG into the training target so
    # the resulting checkpoint can be sampled at guidance_scale=1.0.
    if self._fuse_guidance_scale != 1.0:
        with torch.no_grad():
            noise_pred_uncond = self.student.predict_velocity_with_r(
                noisy,
                t,
                r,
                training_batch,
                conditional=False,
                attn_kind="dense",
            )
        g = float(self._fuse_guidance_scale)
        noise_pred = (noise_pred - (1.0 - g) * noise_pred_uncond) / g

    dF_dt = _central_difference_dF_dt(
        student=self.student,
        batch=training_batch,
        noisy=noisy,
        latents=latents,
        noise=noise,
        t=t,
        r=r,
        delta=self._fd_epsilon,
        num_train_timesteps=n_train,
        attn_kind="dense",
        guidance_scale=self._fuse_guidance_scale,
    )

    # AnyFlow target: target = (eps - x_0) - (t - r) * dF/dt
    # dF/dt is in (velocity per absolute t-unit); (t - r) is in absolute units;
    # the product cancels back to velocity units, matching noise_pred.
    view = [batch_size] + [1] * (latents.ndim - 1)
    target = (noise - latents) - (t - r).view(*view) * dF_dt

    # Per-sample squared error, then per-timestep weight, then scale-balance
    # so the non-diffusion branches stay on the same magnitude as the
    # diffusion branch (matches AnyFlow's stop-grad rescaling).
    per_sample = torch.mean(
        ((noise_pred.float() - target.float())**2).reshape(batch_size, -1),
        dim=-1,
    )
    weight = sched.get_train_weight(t, weight_type=self._weight_type)
    per_sample = per_sample * weight

    with torch.no_grad():
        diff_mask = is_diffusion
        diff_mean = per_sample[diff_mask].mean() if diff_mask.any() else per_sample.mean()
        non_diff_mask = ~diff_mask
        if non_diff_mask.any():
            scale = diff_mean / (per_sample[non_diff_mask] + 1e-5)
        else:
            scale = torch.tensor(1.0, device=device, dtype=per_sample.dtype)
    non_diff_idx = torch.nonzero(non_diff_mask, as_tuple=False).flatten()
    if non_diff_idx.numel() > 0:
        per_sample = per_sample.clone()
        per_sample[non_diff_idx] = per_sample[non_diff_idx] * scale

    total_loss = per_sample.mean()

    loss_map = {"total_loss": total_loss}
    metrics: dict[str, LogScalar] = {
        "diffusion_fraction": float(is_diffusion.float().mean()),
        "consistency_fraction": float(is_consistency.float().mean()),
        "scale_weight_mean": float(scale.mean()) if isinstance(scale, torch.Tensor) else float(scale),
    }
    outputs = {
        "student_ctx": (
            training_batch.timesteps,
            training_batch.attn_metadata,
        ),
    }
    return loss_map, outputs, metrics

def backward(
    self,
    loss_map: dict[str, torch.Tensor],
    outputs: dict[str, Any],
    *,
    grad_accum_rounds: int = 1,
) -> None:
    """Route the loss backward through the student's forward_context
    so attn metadata stays attached during gradient computation."""
    student_ctx = outputs.get("student_ctx")
    if student_ctx is None:
        super().backward(loss_map, outputs, grad_accum_rounds=grad_accum_rounds)
        return
    self.student.backward(
        loss_map["total_loss"],
        student_ctx,
        grad_accum_rounds=grad_accum_rounds,
    )

def _init_optimizer_and_scheduler(self) -> None:
    tc = self.training_config
    params = [p for p in self.student.transformer.parameters() if p.requires_grad]
    (
        self._student_optimizer,
        self._student_lr_scheduler,
    ) = build_optimizer_and_scheduler(
        params=params,
        optimizer_config=tc.optimizer,
        loop_config=tc.loop,
        learning_rate=float(tc.optimizer.learning_rate),
        betas=tc.optimizer.betas,
        scheduler_name=str(tc.optimizer.lr_scheduler),
    )</code></pre></details>
fastvideo.train.methods.distribution_matching.anyflow_pretrain.AnyFlowPretrainMethod.get_optimizers
get_optimizers(iteration: int) -> Sequence[torch.optim.Optimizer]
Source
def get_optimizers(
    self,
    iteration: int,
) -> Sequence[torch.optim.Optimizer]:
    del iteration
    return [self._student_optimizer]
fastvideo.train.methods.distribution_matching.anyflow_pretrain.AnyFlowPretrainMethod.get_lr_schedulers
get_lr_schedulers(iteration: int) -> Sequence[Any]
Source
def get_lr_schedulers(self, iteration: int) -> Sequence[Any]:
    del iteration
    return [self._student_lr_scheduler]
fastvideo.train.methods.distribution_matching.anyflow_pretrain.AnyFlowPretrainMethod.single_train_step
single_train_step(batch: dict[str, Any], iteration: int) -> tuple[dict[str, torch.Tensor], dict[str, Any], dict[str, LogScalar]]
Source
def single_train_step(
    self,
    batch: dict[str, Any],
    iteration: int,
) -> tuple[
        dict[str, torch.Tensor],
        dict[str, Any],
        dict[str, LogScalar],
]:
    del iteration  # AnyFlow pretrain has no iteration-dependent dispatch.
training_batch = self.student.prepare_batch(
    batch,
    generator=self.cuda_generator,
    latents_source=&quot;data&quot;,
)
latents = training_batch.latents  # &#91;B, T, C, H, W&#93; (post-permute in prepare_batch).
if latents is None or latents.ndim != 5:
    raise RuntimeError(&quot;AnyFlow pretrain expects TrainingBatch.latents of shape &quot;
                       &quot;&#91;B, T, C, H, W&#93; after prepare_batch; got &quot;
                       f&quot;&#123;None if latents is None else tuple(latents.shape)&#125;&quot;)
device = latents.device
dtype = latents.dtype
batch_size = int(latents.shape&#91;0&#93;)

# AnyFlow (t, r) sampling — overrides the timestep drawn by
# WanModel._sample_timesteps inside prepare_batch.
t_norm, r_norm, is_diffusion, is_consistency = _sample_pair_timesteps(
    batch_size=batch_size,
    diffusion_ratio=self._diffusion_ratio,
    consistency_ratio=self._consistency_ratio,
    device=device,
    generator=self.cuda_generator,
)

sched = self._flow_map_scheduler
n_train = float(self.student.num_train_timesteps)
t = (sched.apply_shift(t_norm) * n_train).to(device=device, dtype=dtype)
r = (sched.apply_shift(r_norm) * n_train).to(device=device, dtype=dtype)

# Fresh noise drawn from the method&#x27;s RNG; ignore the noise that
# prepare_batch attached (it pairs with the discarded timestep).
noise = torch.randn(
    latents.shape,
    device=device,
    dtype=dtype,
    generator=self.cuda_generator,
)
noisy = sched.add_noise(latents, noise, t)

# Keep training_batch coherent with the new (t, noisy): downstream
# forward_context uses these fields.
training_batch.timesteps = t
training_batch.noise = noise
training_batch.noisy_model_input = noisy.permute(0, 2, 1, 3, 4)

# Student velocity prediction at (t, r).
noise_pred = self.student.predict_velocity_with_r(
    noisy,
    t,
    r,
    training_batch,
    conditional=True,
    attn_kind=&quot;dense&quot;,
)

# Optional guidance distillation — fuse CFG into the training target so
# the resulting checkpoint can be sampled at guidance_scale=1.0.
if self._fuse_guidance_scale != 1.0:
    with torch.no_grad():
        noise_pred_uncond = self.student.predict_velocity_with_r(
            noisy,
            t,
            r,
            training_batch,
            conditional=False,
            attn_kind=&quot;dense&quot;,
        )
    g = float(self._fuse_guidance_scale)
    noise_pred = (noise_pred - (1.0 - g) * noise_pred_uncond) / g

dF_dt = _central_difference_dF_dt(
    student=self.student,
    batch=training_batch,
    noisy=noisy,
    latents=latents,
    noise=noise,
    t=t,
    r=r,
    delta=self._fd_epsilon,
    num_train_timesteps=n_train,
    attn_kind=&quot;dense&quot;,
    guidance_scale=self._fuse_guidance_scale,
)

# AnyFlow target: target = (eps - x_0) - (t - r) * dF/dt
# dF/dt is in (velocity per absolute t-unit); (t - r) is in absolute units;
# the product cancels back to velocity units, matching noise_pred.
view = &#91;batch_size&#93; + &#91;1&#93; * (latents.ndim - 1)
target = (noise - latents) - (t - r).view(*view) * dF_dt

# Per-sample squared error, then per-timestep weight, then scale-balance
# so the non-diffusion branches stay on the same magnitude as the
# diffusion branch (matches AnyFlow&#x27;s stop-grad rescaling).
per_sample = torch.mean(
    ((noise_pred.float() - target.float())**2).reshape(batch_size, -1),
    dim=-1,
)
weight = sched.get_train_weight(t, weight_type=self._weight_type)
per_sample = per_sample * weight

with torch.no_grad():
    diff_mask = is_diffusion
    diff_mean = per_sample&#91;diff_mask&#93;.mean() if diff_mask.any() else per_sample.mean()
    non_diff_mask = ~diff_mask
    if non_diff_mask.any():
        scale = diff_mean / (per_sample&#91;non_diff_mask&#93; + 1e-5)
    else:
        scale = torch.tensor(1.0, device=device, dtype=per_sample.dtype)
non_diff_idx = torch.nonzero(non_diff_mask, as_tuple=False).flatten()
if non_diff_idx.numel() &gt; 0:
    per_sample = per_sample.clone()
    per_sample&#91;non_diff_idx&#93; = per_sample&#91;non_diff_idx&#93; * scale

total_loss = per_sample.mean()

loss_map = &#123;&quot;total_loss&quot;: total_loss&#125;
metrics: dict&#91;str, LogScalar&#93; = &#123;
    &quot;diffusion_fraction&quot;: float(is_diffusion.float().mean()),
    &quot;consistency_fraction&quot;: float(is_consistency.float().mean()),
    &quot;scale_weight_mean&quot;: float(scale.mean()) if isinstance(scale, torch.Tensor) else float(scale),
&#125;
outputs = &#123;
    &quot;student_ctx&quot;: (
        training_batch.timesteps,
        training_batch.attn_metadata,
    ),
&#125;
return loss_map, outputs, metrics</code></pre></details>
fastvideo.train.methods.distribution_matching.anyflow_pretrain.AnyFlowPretrainMethod.backward
backward(loss_map: dict[str, torch.Tensor], outputs: dict[str, Any], *, grad_accum_rounds: int = 1) -> None
Route the loss backward through the student's forward_context
so attn metadata stays attached during gradient computation.
Source
def backward(
    self,
    loss_map: dict[str, torch.Tensor],
    outputs: dict[str, Any],
    *,
    grad_accum_rounds: int = 1,
) -> None:
    """Route the loss backward through the student's forward_context
    so attn metadata stays attached during gradient computation."""
    student_ctx = outputs.get("student_ctx")
    if student_ctx is None:
        super().backward(loss_map, outputs, grad_accum_rounds=grad_accum_rounds)
        return
    self.student.backward(
        loss_map["total_loss"],
        student_ctx,
        grad_accum_rounds=grad_accum_rounds,
    )