577 lines
21 KiB
Python
577 lines
21 KiB
Python
# SPDX-License-Identifier: Apache-2.0
|
|
"""Self-Forcing distillation method (algorithm layer)."""
|
|
|
|
from __future__ import annotations
|
|
|
|
from typing import Any, Literal, TYPE_CHECKING
|
|
|
|
import torch
|
|
import torch.distributed as dist
|
|
|
|
from fastvideo.train.models.base import (
|
|
CausalModelBase,
|
|
ModelBase,
|
|
)
|
|
from fastvideo.train.methods.distribution_matching.dmd2 import (
|
|
DMD2Method, )
|
|
from fastvideo.train.utils.config import (
|
|
get_optional_float,
|
|
get_optional_int,
|
|
)
|
|
from fastvideo.models.schedulers.scheduling_self_forcing_flow_match import (
|
|
SelfForcingFlowMatchScheduler, )
|
|
from fastvideo.models.utils import pred_noise_to_pred_video
|
|
|
|
if TYPE_CHECKING:
|
|
from fastvideo.pipelines import TrainingBatch
|
|
|
|
|
|
def _require_bool(raw: Any, *, where: str) -> bool:
|
|
if isinstance(raw, bool):
|
|
return raw
|
|
raise ValueError(f"Expected bool at {where}, got {type(raw).__name__}")
|
|
|
|
|
|
def _require_str(raw: Any, *, where: str) -> str:
|
|
if not isinstance(raw, str) or not raw.strip():
|
|
raise ValueError(f"Expected non-empty string at {where}")
|
|
return raw
|
|
|
|
|
|
class SelfForcingMethod(DMD2Method):
|
|
"""Self-Forcing DMD2 (distribution matching) method.
|
|
|
|
Requires a causal student implementing ``CausalModelBase``.
|
|
"""
|
|
|
|
def __init__(
|
|
self,
|
|
*,
|
|
cfg: Any,
|
|
role_models: dict[str, ModelBase],
|
|
) -> None:
|
|
super().__init__(
|
|
cfg=cfg,
|
|
role_models=role_models,
|
|
)
|
|
|
|
# Validate causal student.
|
|
if not isinstance(self.student, CausalModelBase):
|
|
raise ValueError("SelfForcingMethod requires a causal student "
|
|
"implementing CausalModelBase.")
|
|
|
|
if self._rollout_mode != "simulate":
|
|
raise ValueError("SelfForcingMethod only supports "
|
|
"method_config.rollout_mode='simulate'")
|
|
|
|
mcfg = self.method_config
|
|
|
|
chunk_size = get_optional_int(
|
|
mcfg,
|
|
"chunk_size",
|
|
where="method_config.chunk_size",
|
|
)
|
|
if chunk_size is None:
|
|
chunk_size = 3
|
|
if chunk_size <= 0:
|
|
raise ValueError("method_config.chunk_size must be a positive "
|
|
f"integer, got {chunk_size}")
|
|
self._chunk_size = int(chunk_size)
|
|
|
|
sample_type_raw = mcfg.get("student_sample_type", "sde")
|
|
sample_type = _require_str(
|
|
sample_type_raw,
|
|
where="method_config.student_sample_type",
|
|
)
|
|
sample_type = sample_type.strip().lower()
|
|
if sample_type not in {"sde", "ode"}:
|
|
raise ValueError("method_config.student_sample_type must be one "
|
|
f"of {{sde, ode}}, got {sample_type_raw!r}")
|
|
self._student_sample_type: Literal["sde", "ode"] = (
|
|
sample_type # type: ignore[assignment]
|
|
)
|
|
|
|
same_step_raw = mcfg.get("same_step_across_blocks", False)
|
|
if same_step_raw is None:
|
|
same_step_raw = False
|
|
self._same_step_across_blocks = _require_bool(
|
|
same_step_raw,
|
|
where="method_config.same_step_across_blocks",
|
|
)
|
|
|
|
last_step_raw = mcfg.get("last_step_only", False)
|
|
if last_step_raw is None:
|
|
last_step_raw = False
|
|
self._last_step_only = _require_bool(
|
|
last_step_raw,
|
|
where="method_config.last_step_only",
|
|
)
|
|
|
|
context_noise = get_optional_float(
|
|
mcfg,
|
|
"context_noise",
|
|
where="method_config.context_noise",
|
|
)
|
|
if context_noise is None:
|
|
context_noise = 0.0
|
|
if context_noise < 0.0:
|
|
raise ValueError("method_config.context_noise must be >= 0, "
|
|
f"got {context_noise}")
|
|
self._context_noise = float(context_noise)
|
|
|
|
enable_grad_raw = mcfg.get("enable_gradient_in_rollout", True)
|
|
if enable_grad_raw is None:
|
|
enable_grad_raw = True
|
|
self._enable_gradient_in_rollout = _require_bool(
|
|
enable_grad_raw,
|
|
where="method_config.enable_gradient_in_rollout",
|
|
)
|
|
|
|
start_grad_frame = get_optional_int(
|
|
mcfg,
|
|
"start_gradient_frame",
|
|
where="method_config.start_gradient_frame",
|
|
)
|
|
if start_grad_frame is None:
|
|
start_grad_frame = 0
|
|
if start_grad_frame < 0:
|
|
raise ValueError("method_config.start_gradient_frame must be "
|
|
f">= 0, got {start_grad_frame}")
|
|
self._start_gradient_frame = int(start_grad_frame)
|
|
|
|
shift = float(getattr(
|
|
self.training_config.pipeline_config,
|
|
"flow_shift",
|
|
0.0,
|
|
) or 0.0)
|
|
self._sf_scheduler = SelfForcingFlowMatchScheduler(
|
|
num_inference_steps=1000,
|
|
num_train_timesteps=int(self.student.num_train_timesteps),
|
|
shift=shift,
|
|
sigma_min=0.0,
|
|
extra_one_step=True,
|
|
training=True,
|
|
)
|
|
|
|
self._sf_denoising_step_list: torch.Tensor | None = None
|
|
|
|
def _get_denoising_step_list(self, device: torch.device) -> torch.Tensor:
|
|
if (self._sf_denoising_step_list is not None and self._sf_denoising_step_list.device == device):
|
|
return self._sf_denoising_step_list
|
|
|
|
raw = self.method_config.get("dmd_denoising_steps", None)
|
|
if not isinstance(raw, list) or not raw:
|
|
raise ValueError("method_config.dmd_denoising_steps must be set "
|
|
"for self_forcing")
|
|
steps = torch.tensor(
|
|
[int(s) for s in raw],
|
|
dtype=torch.long,
|
|
device=device,
|
|
)
|
|
|
|
warp = self.method_config.get("warp_denoising_step", None)
|
|
if warp is None:
|
|
warp = False
|
|
if bool(warp):
|
|
timesteps = torch.cat((
|
|
self._sf_scheduler.timesteps.to("cpu"),
|
|
torch.tensor([0], dtype=torch.float32),
|
|
)).to(device)
|
|
steps = timesteps[int(self.student.num_train_timesteps) - steps]
|
|
|
|
self._sf_denoising_step_list = steps
|
|
return steps
|
|
|
|
def _predict_x0_with_scheduler(
|
|
self,
|
|
model: ModelBase,
|
|
noisy_latents: torch.Tensor,
|
|
timestep: torch.Tensor,
|
|
batch: TrainingBatch,
|
|
*,
|
|
conditional: bool,
|
|
attn_kind: Literal["dense", "vsa"],
|
|
) -> torch.Tensor:
|
|
pred_noise = model.predict_noise(
|
|
noisy_latents,
|
|
timestep,
|
|
batch,
|
|
conditional=conditional,
|
|
cfg_uncond=self._cfg_uncond,
|
|
attn_kind=attn_kind,
|
|
)
|
|
pred_x0 = pred_noise_to_pred_video(
|
|
pred_noise=pred_noise.flatten(0, 1),
|
|
noise_input_latent=noisy_latents.flatten(0, 1),
|
|
timestep=timestep,
|
|
scheduler=self._sf_scheduler,
|
|
).unflatten(0, pred_noise.shape[:2])
|
|
return pred_x0
|
|
|
|
def _sf_add_noise(
|
|
self,
|
|
clean_latents: torch.Tensor,
|
|
noise: torch.Tensor,
|
|
timestep: torch.Tensor,
|
|
) -> torch.Tensor:
|
|
b, t = clean_latents.shape[:2]
|
|
noisy = self._sf_scheduler.add_noise(
|
|
clean_latents.flatten(0, 1),
|
|
noise.flatten(0, 1),
|
|
timestep,
|
|
).unflatten(0, (b, t))
|
|
return noisy
|
|
|
|
def _timestep_to_sigma(self, timestep: torch.Tensor) -> torch.Tensor:
|
|
sigmas = self._sf_scheduler.sigmas.to(device=timestep.device, dtype=torch.float32)
|
|
timesteps = self._sf_scheduler.timesteps.to(device=timestep.device, dtype=torch.float32)
|
|
t = timestep.to(device=timestep.device, dtype=torch.float32)
|
|
if t.ndim == 2:
|
|
t = t.flatten(0, 1)
|
|
elif t.ndim == 1 and t.numel() == 1:
|
|
t = t.expand(1)
|
|
elif t.ndim != 1:
|
|
raise ValueError("Invalid timestep shape: "
|
|
f"{tuple(timestep.shape)}")
|
|
idx = torch.argmin(
|
|
(timesteps.unsqueeze(0) - t.unsqueeze(1)).abs(),
|
|
dim=1,
|
|
)
|
|
return sigmas[idx]
|
|
|
|
def _sample_exit_indices(
|
|
self,
|
|
*,
|
|
num_blocks: int,
|
|
num_steps: int,
|
|
device: torch.device,
|
|
) -> list[int]:
|
|
if num_blocks <= 0:
|
|
return []
|
|
if num_steps <= 0:
|
|
raise ValueError("num_steps must be positive")
|
|
|
|
shape = ((1, ) if self._same_step_across_blocks else (num_blocks, ))
|
|
|
|
if not dist.is_initialized() or dist.get_rank() == 0:
|
|
if self._last_step_only:
|
|
indices = torch.full(
|
|
shape,
|
|
num_steps - 1,
|
|
dtype=torch.long,
|
|
device=device,
|
|
)
|
|
else:
|
|
indices = torch.randint(
|
|
low=0,
|
|
high=num_steps,
|
|
size=shape,
|
|
device=device,
|
|
generator=self.cuda_generator,
|
|
)
|
|
else:
|
|
indices = torch.empty(shape, dtype=torch.long, device=device)
|
|
|
|
if dist.is_initialized():
|
|
dist.broadcast(indices, src=0)
|
|
|
|
if self._same_step_across_blocks:
|
|
return [int(indices.item()) for _ in range(num_blocks)]
|
|
return [int(i) for i in indices.tolist()]
|
|
|
|
def _student_rollout(self, batch: Any, *, with_grad: bool) -> torch.Tensor:
|
|
if not isinstance(self.student, CausalModelBase):
|
|
raise ValueError("SelfForcingMethod requires a causal student "
|
|
"implementing CausalModelBase.")
|
|
return self._student_rollout_streaming(batch, with_grad=with_grad)
|
|
|
|
def _student_rollout_streaming(self, batch: Any, *, with_grad: bool) -> torch.Tensor:
|
|
assert isinstance(self.student, CausalModelBase)
|
|
latents = batch.latents
|
|
if latents is None:
|
|
raise RuntimeError("TrainingBatch.latents is required for "
|
|
"self-forcing rollout")
|
|
if latents.ndim != 5:
|
|
raise ValueError("TrainingBatch.latents must be [B, T, C, H, W]"
|
|
f", got shape={tuple(latents.shape)}")
|
|
|
|
device = latents.device
|
|
dtype = latents.dtype
|
|
batch_size = int(latents.shape[0])
|
|
num_frames = int(latents.shape[1])
|
|
|
|
denoising_steps = self._get_denoising_step_list(device)
|
|
num_steps = int(denoising_steps.numel())
|
|
|
|
noise_full = torch.randn_like(latents, device=device, dtype=dtype)
|
|
|
|
chunk = int(self._chunk_size)
|
|
if chunk <= 0:
|
|
raise ValueError("chunk_size must be positive")
|
|
|
|
remaining = num_frames % chunk
|
|
num_blocks = num_frames // chunk
|
|
if num_blocks == 0:
|
|
num_blocks = 1
|
|
remaining = num_frames
|
|
|
|
exit_indices = self._sample_exit_indices(
|
|
num_blocks=num_blocks,
|
|
num_steps=num_steps,
|
|
device=device,
|
|
)
|
|
|
|
denoised_blocks: list[torch.Tensor] = []
|
|
|
|
cache_tag = "pos"
|
|
self.student.clear_caches(cache_tag=cache_tag)
|
|
|
|
for block_idx in range(num_blocks):
|
|
if block_idx == 0:
|
|
start = 0
|
|
end = remaining + chunk if remaining else chunk
|
|
else:
|
|
start = remaining + block_idx * chunk
|
|
end = remaining + (block_idx + 1) * chunk
|
|
start = int(start)
|
|
end = int(min(end, num_frames))
|
|
if start >= end:
|
|
break
|
|
|
|
noisy_block = noise_full[:, start:end]
|
|
exit_idx = int(exit_indices[block_idx])
|
|
|
|
for step_idx, current_timestep in enumerate(denoising_steps):
|
|
exit_flag = step_idx == exit_idx
|
|
|
|
timestep_block = (current_timestep * torch.ones(
|
|
(batch_size, end - start),
|
|
device=device,
|
|
dtype=torch.float32,
|
|
))
|
|
|
|
enable_grad = (bool(with_grad) and bool(self._enable_gradient_in_rollout) and torch.is_grad_enabled()
|
|
and start >= int(self._start_gradient_frame))
|
|
|
|
if not exit_flag:
|
|
with torch.no_grad():
|
|
pred_noise = (self.student.predict_noise_streaming(
|
|
noisy_block,
|
|
timestep_block,
|
|
batch,
|
|
conditional=True,
|
|
cache_tag=cache_tag,
|
|
store_kv=False,
|
|
cur_start_frame=start,
|
|
cfg_uncond=self._cfg_uncond,
|
|
attn_kind="vsa",
|
|
))
|
|
if pred_noise is None:
|
|
raise RuntimeError("predict_noise_streaming "
|
|
"returned None "
|
|
"(store_kv=False)")
|
|
pred_x0_chunk = pred_noise_to_pred_video(
|
|
pred_noise=pred_noise.flatten(0, 1),
|
|
noise_input_latent=(noisy_block.flatten(0, 1)),
|
|
timestep=timestep_block,
|
|
scheduler=self._sf_scheduler,
|
|
).unflatten(0, pred_noise.shape[:2])
|
|
|
|
if step_idx + 1 >= num_steps:
|
|
break
|
|
next_timestep = denoising_steps[step_idx + 1]
|
|
if self._student_sample_type == "sde":
|
|
noisy_block = self._sf_add_noise(
|
|
pred_x0_chunk,
|
|
torch.randn_like(pred_x0_chunk),
|
|
next_timestep * torch.ones(
|
|
(batch_size, end - start),
|
|
device=device,
|
|
dtype=torch.float32,
|
|
),
|
|
)
|
|
else:
|
|
sigma_cur = self._timestep_to_sigma(timestep_block).view(batch_size, end - start, 1, 1, 1)
|
|
sigma_next = self._timestep_to_sigma(next_timestep * torch.ones(
|
|
(batch_size, end - start),
|
|
device=device,
|
|
dtype=torch.float32,
|
|
)).view(batch_size, end - start, 1, 1, 1)
|
|
eps = (noisy_block - (1 - sigma_cur) * pred_x0_chunk) / sigma_cur.clamp_min(1e-8)
|
|
noisy_block = ((1 - sigma_next) * pred_x0_chunk + sigma_next * eps)
|
|
continue
|
|
|
|
with torch.set_grad_enabled(enable_grad):
|
|
pred_noise = (self.student.predict_noise_streaming(
|
|
noisy_block,
|
|
timestep_block,
|
|
batch,
|
|
conditional=True,
|
|
cache_tag=cache_tag,
|
|
store_kv=False,
|
|
cur_start_frame=start,
|
|
cfg_uncond=self._cfg_uncond,
|
|
attn_kind="vsa",
|
|
))
|
|
if pred_noise is None:
|
|
raise RuntimeError("predict_noise_streaming returned "
|
|
"None (store_kv=False)")
|
|
pred_x0_chunk = pred_noise_to_pred_video(
|
|
pred_noise=pred_noise.flatten(0, 1),
|
|
noise_input_latent=(noisy_block.flatten(0, 1)),
|
|
timestep=timestep_block,
|
|
scheduler=self._sf_scheduler,
|
|
).unflatten(0, pred_noise.shape[:2])
|
|
break
|
|
|
|
denoised_blocks.append(pred_x0_chunk)
|
|
|
|
with torch.no_grad():
|
|
if self._context_noise > 0.0:
|
|
context_timestep = torch.ones(
|
|
(batch_size, end - start),
|
|
device=device,
|
|
dtype=torch.float32,
|
|
) * float(self._context_noise)
|
|
context_latents = self._sf_add_noise(
|
|
pred_x0_chunk.detach(),
|
|
torch.randn_like(pred_x0_chunk),
|
|
context_timestep,
|
|
)
|
|
else:
|
|
context_timestep = torch.zeros(
|
|
(batch_size, end - start),
|
|
device=device,
|
|
dtype=torch.float32,
|
|
)
|
|
context_latents = pred_x0_chunk.detach()
|
|
|
|
_ = self.student.predict_noise_streaming(
|
|
context_latents,
|
|
context_timestep,
|
|
batch,
|
|
conditional=True,
|
|
cache_tag=cache_tag,
|
|
store_kv=True,
|
|
cur_start_frame=start,
|
|
cfg_uncond=self._cfg_uncond,
|
|
attn_kind="vsa",
|
|
)
|
|
|
|
if not denoised_blocks:
|
|
raise RuntimeError("Self-forcing rollout produced no blocks")
|
|
|
|
self.student.clear_caches(cache_tag=cache_tag)
|
|
return torch.cat(denoised_blocks, dim=1)
|
|
|
|
def _critic_flow_matching_loss(self, batch: Any) -> tuple[torch.Tensor, Any, dict[str, Any]]:
|
|
with torch.no_grad():
|
|
generator_pred_x0 = self._student_rollout(batch, with_grad=False)
|
|
|
|
device = generator_pred_x0.device
|
|
fake_score_timestep = torch.randint(
|
|
0,
|
|
int(self.student.num_train_timesteps),
|
|
[1],
|
|
device=device,
|
|
dtype=torch.long,
|
|
generator=self.cuda_generator,
|
|
)
|
|
fake_score_timestep = (self.student.shift_and_clamp_timestep(fake_score_timestep))
|
|
|
|
noise = torch.randn(
|
|
generator_pred_x0.shape,
|
|
device=device,
|
|
dtype=generator_pred_x0.dtype,
|
|
generator=self.cuda_generator,
|
|
)
|
|
noisy_x0 = self._sf_add_noise(generator_pred_x0, noise, fake_score_timestep)
|
|
|
|
pred_noise = self.critic.predict_noise(
|
|
noisy_x0,
|
|
fake_score_timestep,
|
|
batch,
|
|
conditional=True,
|
|
cfg_uncond=self._cfg_uncond,
|
|
attn_kind="dense",
|
|
)
|
|
target = noise - generator_pred_x0
|
|
flow_matching_loss = torch.mean((pred_noise - target)**2)
|
|
|
|
batch.fake_score_latent_vis_dict = {
|
|
"generator_pred_video": generator_pred_x0,
|
|
"fake_score_timestep": fake_score_timestep,
|
|
}
|
|
outputs = {"fake_score_latent_vis_dict": (batch.fake_score_latent_vis_dict)}
|
|
return (
|
|
flow_matching_loss,
|
|
(batch.timesteps, batch.attn_metadata),
|
|
outputs,
|
|
)
|
|
|
|
def _dmd_loss(
|
|
self,
|
|
generator_pred_x0: torch.Tensor,
|
|
batch: Any,
|
|
) -> torch.Tensor:
|
|
guidance_scale = get_optional_float(
|
|
self.method_config,
|
|
"real_score_guidance_scale",
|
|
where="method.real_score_guidance_scale",
|
|
)
|
|
if guidance_scale is None:
|
|
guidance_scale = 1.0
|
|
device = generator_pred_x0.device
|
|
|
|
with torch.no_grad():
|
|
timestep = torch.randint(
|
|
0,
|
|
int(self.student.num_train_timesteps),
|
|
[1],
|
|
device=device,
|
|
dtype=torch.long,
|
|
generator=self.cuda_generator,
|
|
)
|
|
timestep = self.student.shift_and_clamp_timestep(timestep)
|
|
|
|
noise = torch.randn(
|
|
generator_pred_x0.shape,
|
|
device=device,
|
|
dtype=generator_pred_x0.dtype,
|
|
generator=self.cuda_generator,
|
|
)
|
|
noisy_latents = self._sf_add_noise(generator_pred_x0, noise, timestep)
|
|
|
|
faker_x0 = self._predict_x0_with_scheduler(
|
|
self.critic,
|
|
noisy_latents,
|
|
timestep,
|
|
batch,
|
|
conditional=True,
|
|
attn_kind="dense",
|
|
)
|
|
real_cond_x0 = self._predict_x0_with_scheduler(
|
|
self.teacher,
|
|
noisy_latents,
|
|
timestep,
|
|
batch,
|
|
conditional=True,
|
|
attn_kind="dense",
|
|
)
|
|
real_uncond_x0 = self._predict_x0_with_scheduler(
|
|
self.teacher,
|
|
noisy_latents,
|
|
timestep,
|
|
batch,
|
|
conditional=False,
|
|
attn_kind="dense",
|
|
)
|
|
real_cfg_x0 = real_uncond_x0 + (real_cond_x0 - real_uncond_x0) * guidance_scale
|
|
|
|
denom = torch.abs(generator_pred_x0 - real_cfg_x0).mean()
|
|
grad = (faker_x0 - real_cfg_x0) / denom
|
|
grad = torch.nan_to_num(grad)
|
|
|
|
loss = 0.5 * torch.mean((generator_pred_x0.float() - (generator_pred_x0.float() - grad.float()).detach())**2)
|
|
return loss
|