streaming_long_tuning
LongLive-style multi-stage self-forcing for streaming rollouts.
fastvideo.train.methods.distribution_matching.streaming_long_tuning.init_logger
init_logger(name: str) -> _FastvideoLoggerThe main purpose of this function is to ensure that loggers are retrieved in such a way that we can be sure the root fastvideo logger has already been configured.
fastvideo.train.methods.distribution_matching.streaming_long_tuning.pred_noise_to_pred_video
pred_noise_to_pred_video(pred_noise: torch.Tensor, noise_input_latent: torch.Tensor, timestep: torch.Tensor, scheduler: Any) -> torch.TensorConvert predicted noise to clean latent.
pred_noise: the predicted noise with shape [B, C, H, W] where B is batch_size or batch_size * num_frames noise_input_latent: the noisy latent with shape [B, C, H, W], timestep: the timestep with shape [1] or [bs * num_frames] or [bs, num_frames] scheduler: the scheduler
Returns
torch.Tensorthe predicted video with shape [B, C, H, W]
fastvideo.train.methods.distribution_matching.streaming_long_tuning.SelfForcingMethod
SelfForcingMethod(*, cfg: Any, role_models: dict[str, ModelBase])Self-Forcing DMD2 (distribution matching) method.
Requires a causal student implementing CausalModelBase.
fastvideo.train.methods.distribution_matching.streaming_long_tuning.CausalModelBase
CausalModelBase(*, trainable: bool = True, lora: LoraConfig | dict[str, Any] | None = None, attention_backend: AttentionBackendEnum | str | None = None)Extension for causal / streaming model plugins.
Cache state is internal to the model instance and keyed by cache_tag (no role handle needed).
fastvideo.train.methods.distribution_matching.streaming_long_tuning.CausalModelBase.clear_caches
clear_caches(*, cache_tag: str = 'pos') -> NoneClear internal caches before starting a new rollout.
fastvideo.train.methods.distribution_matching.streaming_long_tuning.CausalModelBase.predict_noise_streaming
predict_noise_streaming(noisy_latents: torch.Tensor, timestep: torch.Tensor, batch: TrainingBatch, *, conditional: bool, cache_tag: str = 'pos', store_kv: bool = False, cur_start_frame: int = 0, cfg_uncond: dict[str, Any] | None = None, attn_kind: Literal['dense', 'vsa'] = 'dense') -> torch.Tensor | NoneStreaming predict-noise that may update internal caches.
fastvideo.train.methods.distribution_matching.streaming_long_tuning.CausalModelBase.predict_x0_streaming
predict_x0_streaming(noisy_latents: torch.Tensor, timestep: torch.Tensor, batch: TrainingBatch, *, conditional: bool, cache_tag: str = 'pos', store_kv: bool = False, cur_start_frame: int = 0, cfg_uncond: dict[str, Any] | None = None, attn_kind: Literal['dense', 'vsa'] = 'dense') -> torch.Tensor | NonePredict x0 streaming via
predict_noise_streaming + conversion.
fastvideo.train.methods.distribution_matching.streaming_long_tuning.get_optional_float
get_optional_float(mapping: dict[str, Any], key: str, *, where: str) -> float | Nonefastvideo.train.methods.distribution_matching.streaming_long_tuning.get_optional_int
get_optional_int(mapping: dict[str, Any], key: str, *, where: str) -> int | Nonefastvideo.train.methods.distribution_matching.streaming_long_tuning.DistillStage
DistillStage(name: str, start_step: int, end_step: int | None, num_latent_t: int, streaming_training: bool, streaming_chunk_size: int | None = None, streaming_max_length: int | None = None, streaming_min_new_frame: int | None = None, streaming_fixed_overlap_latents: int | None = None)Resolved multi-stage distillation stage.
Source
@dataclass(frozen=True, slots=True)
class DistillStage:
"""Resolved multi-stage distillation stage."""
name: str
start_step: int
end_step: int | None
num_latent_t: int
streaming_training: bool
streaming_chunk_size: int | None = None
streaming_max_length: int | None = None
streaming_min_new_frame: int | None = None
streaming_fixed_overlap_latents: int | None = None</code></pre></details>
fastvideo.train.methods.distribution_matching.streaming_long_tuning.parse_multi_phased_distill_schedule
parse_multi_phased_distill_schedule(raw: Any, *, default_num_latent_t: int, default_streaming_chunk_size: int | None = None, default_streaming_max_length: int | None = None) -> list[DistillStage]
Parse a compact/list multi-stage schedule.
Preferred YAML form:
[{stage: self_forcing, end_step: 700, num_latent_t: 21}, ...]
Compact string form accepts "700:21,3000:240" and treats the
first stage as non-streaming self-forcing and later stages as streaming.
Source
def parse_multi_phased_distill_schedule(
raw: Any,
*,
default_num_latent_t: int,
default_streaming_chunk_size: int | None = None,
default_streaming_max_length: int | None = None,
) -> list[DistillStage]:
"""Parse a compact/list multi-stage schedule.
Preferred YAML form:
``[{stage: self_forcing, end_step: 700, num_latent_t: 21}, ...]``
Compact string form accepts ``"700:21,3000:240"`` and treats the
first stage as non-streaming self-forcing and later stages as streaming.
"""
if raw is None or raw == "":
max_length = default_streaming_max_length or default_num_latent_t
return [
DistillStage(
name="streaming_long",
start_step=0,
end_step=None,
num_latent_t=int(max_length),
streaming_training=True,
streaming_chunk_size=default_streaming_chunk_size,
streaming_max_length=int(max_length),
)
]
stages: list[DistillStage] = []
previous_end = 0
if isinstance(raw, str):
parts = [part.strip() for part in raw.split(",") if part.strip()]
for idx, part in enumerate(parts):
fields = [field.strip() for field in part.split(":")]
end_step: int | None
if len(fields) == 2:
start_step = previous_end
end_step = _as_int(
fields[0],
where="multi_phased_distill_schedule.end_step",
)
num_latent_t = _as_int(
fields[1],
where="multi_phased_distill_schedule.num_latent_t",
)
elif len(fields) == 3:
start_step = _as_int(
fields[0],
where="multi_phased_distill_schedule.start_step",
)
end_step = _as_int(
fields[1],
where="multi_phased_distill_schedule.end_step",
)
num_latent_t = _as_int(
fields[2],
where="multi_phased_distill_schedule.num_latent_t",
)
else:
raise ValueError("multi_phased_distill_schedule string entries must be "
"'end_step:num_latent_t' or 'start:end:num_latent_t', "
f"got {part!r}")
streaming = idx > 0
stages.append(
DistillStage(
name="streaming_long" if streaming else "self_forcing",
start_step=start_step,
end_step=end_step,
num_latent_t=num_latent_t,
streaming_training=streaming,
streaming_chunk_size=(default_streaming_chunk_size if streaming else None),
streaming_max_length=num_latent_t if streaming else None,
))
previous_end = end_step
elif isinstance(raw, list | tuple):
for idx, stage_raw in enumerate(raw):
if not isinstance(stage_raw, dict):
raise ValueError("multi_phased_distill_schedule list entries must be dicts")
name = str(stage_raw.get("stage", "") or stage_raw.get("name", "")).strip()
streaming_raw = stage_raw.get("streaming_training", None)
if streaming_raw is None:
streaming = name in {"streaming_long", "long", "streaming"}
else:
streaming = _as_bool(
streaming_raw,
where=f"multi_phased_distill_schedule[{idx}].streaming_training",
)
if not name:
name = "streaming_long" if streaming else "self_forcing"
start_step = _as_int(
stage_raw.get("start_step", previous_end),
where=f"multi_phased_distill_schedule[{idx}].start_step",
)
end_raw = stage_raw.get("end_step", None)
end_step = (None if end_raw is None else _as_int(
end_raw,
where=f"multi_phased_distill_schedule[{idx}].end_step",
))
num_latent_t = _as_int(
stage_raw.get(
"num_latent_t",
stage_raw.get(
"streaming_max_length",
stage_raw.get("max_length", default_num_latent_t),
),
),
where=f"multi_phased_distill_schedule[{idx}].num_latent_t",
)
streaming_max = stage_raw.get("streaming_max_length", None)
chunk_size = stage_raw.get("streaming_chunk_size", None)
min_new = stage_raw.get("streaming_min_new_frame", None)
fixed_overlap = stage_raw.get("streaming_fixed_overlap_latents", None)
stages.append(
DistillStage(
name=name,
start_step=start_step,
end_step=end_step,
num_latent_t=num_latent_t,
streaming_training=streaming,
streaming_chunk_size=(None if chunk_size is None else _as_int(
chunk_size,
where=("multi_phased_distill_schedule"
f"[{idx}].streaming_chunk_size"),
)),
streaming_max_length=(None if streaming_max is None else _as_int(
streaming_max,
where=("multi_phased_distill_schedule"
f"[{idx}].streaming_max_length"),
)),
streaming_min_new_frame=(None if min_new is None else _as_int(
min_new,
where=("multi_phased_distill_schedule"
f"[{idx}].streaming_min_new_frame"),
)),
streaming_fixed_overlap_latents=(None if fixed_overlap is None else _as_int(
fixed_overlap,
where=("multi_phased_distill_schedule"
f"[{idx}].streaming_fixed_overlap_latents"),
)),
))
if end_step is not None:
previous_end = end_step
else:
raise ValueError("multi_phased_distill_schedule must be a list, string, or empty")
if not stages:
raise ValueError("multi_phased_distill_schedule produced no stages")
previous_end = 0
for stage in stages:
if stage.start_step < previous_end:
raise ValueError("Distillation stages must be ordered and non-overlapping")
if stage.end_step is not None and stage.end_step <= stage.start_step:
raise ValueError("Distillation stage end_step must be > start_step")
if stage.num_latent_t <= 0:
raise ValueError("Distillation stage num_latent_t must be positive")
if stage.streaming_training:
max_length = stage.streaming_max_length or stage.num_latent_t
chunk_size = stage.streaming_chunk_size or default_streaming_chunk_size
if max_length <= 0:
raise ValueError("streaming_max_length must be positive")
if chunk_size is None or chunk_size <= 0:
raise ValueError("streaming_chunk_size must be positive")
if stage.streaming_fixed_overlap_latents is not None:
if stage.streaming_fixed_overlap_latents < 0:
raise ValueError("streaming_fixed_overlap_latents must be non-negative")
if stage.streaming_fixed_overlap_latents >= chunk_size:
raise ValueError("streaming_fixed_overlap_latents must be smaller than "
"streaming_chunk_size")
if stage.end_step is not None:
previous_end = stage.end_step
return stages</code></pre></details>
fastvideo.train.methods.distribution_matching.streaming_long_tuning.select_distill_stage
select_distill_stage(stages: list[DistillStage], iteration: int) -> DistillStage
Select the active stage for iteration.
Source
def select_distill_stage(
stages: list[DistillStage],
iteration: int,
) -> DistillStage:
"""Select the active stage for ``iteration``."""
step = int(iteration)
for stage in stages:
if stage.end_step is None:
if step >= stage.start_step:
return stage
elif stage.start_step <= step < stage.end_step:
return stage
return stages[-1]</code></pre></details>
fastvideo.train.methods.distribution_matching.streaming_long_tuning.StreamingLongTuningMethod
StreamingLongTuningMethod(*, cfg: Any, role_models: dict[str, Any])
Two-stage MatrixGame/LongLive-style self-forcing method.
Stage 1 uses the existing full self-forcing rollout over a short latent
horizon. Stage 2 uses a persistent streaming sequence and trains on chunks
generated with the causal student cache.
Source
class StreamingLongTuningMethod(SelfForcingMethod):
"""Two-stage MatrixGame/LongLive-style self-forcing method.
Stage 1 uses the existing full self-forcing rollout over a short latent
horizon. Stage 2 uses a persistent streaming sequence and trains on chunks
generated with the causal student cache.
"""
def __init__(
self,
*,
cfg: Any,
role_models: dict[str, Any],
) -> None:
pipeline_config = getattr(cfg.training, "pipeline_config", None)
vae_config = getattr(pipeline_config, "vae_config", None)
if vae_config is not None:
vae_config.load_encoder = True
vae_config.load_decoder = True
super().__init__(cfg=cfg, role_models=role_models)
if not isinstance(self.student, CausalModelBase):
raise ValueError("StreamingLongTuningMethod requires a causal student")
mcfg = self.method_config
default_chunk = get_optional_int(
mcfg,
"streaming_chunk_size",
where="method.streaming_chunk_size",
)
default_max = get_optional_int(
mcfg,
"streaming_max_length",
where="method.streaming_max_length",
)
default_num_latent_t = int(self.training_config.data.num_latent_t)
self._distill_stages = parse_multi_phased_distill_schedule(
mcfg.get("multi_phased_distill_schedule", None),
default_num_latent_t=default_num_latent_t,
default_streaming_chunk_size=default_chunk,
default_streaming_max_length=default_max,
)
self._streaming_state: _StreamingState | None = None
self._streaming_anchor_inject_k = int(
get_optional_int(
mcfg,
"streaming_anchor_inject_k",
where="method.streaming_anchor_inject_k",
) or 1)
if self._streaming_anchor_inject_k < 1:
raise ValueError("method.streaming_anchor_inject_k must be >= 1")
reencode_anchor_raw = mcfg.get("streaming_reencode_overlap_anchor", True)
if reencode_anchor_raw is None:
reencode_anchor_raw = True
self._streaming_reencode_overlap_anchor = _as_bool(
reencode_anchor_raw,
where="method.streaming_reencode_overlap_anchor",
)
require_full_blocks_raw = mcfg.get("streaming_require_full_blocks", False)
if require_full_blocks_raw is None:
require_full_blocks_raw = False
self._streaming_require_full_blocks = _as_bool(
require_full_blocks_raw,
where="method.streaming_require_full_blocks",
)
real_score_cfg_raw = mcfg.get("real_score_cfg", True)
if real_score_cfg_raw is None:
real_score_cfg_raw = True
self._real_score_cfg = _as_bool(
real_score_cfg_raw,
where="method.real_score_cfg",
)
self._score_min_timestep = self._score_timestep_from_ratio("min_timestep_ratio")
self._score_max_timestep = self._score_timestep_from_ratio("max_timestep_ratio")
max_stage_latents = max(s.num_latent_t for s in self._distill_stages)
if max_stage_latents > default_num_latent_t:
raise ValueError("training.data.num_latent_t must be at least the largest "
"multi-stage horizon so prepared conditioning is not trimmed "
"before streaming: "
f"num_latent_t={default_num_latent_t}, "
f"max_stage_latents={max_stage_latents}")
def single_train_step(
self,
batch: dict[str, Any],
iteration: int,
) -> tuple[
dict[str, torch.Tensor],
dict[str, Any],
dict[str, LogScalar],
]:
stage = select_distill_stage(self._distill_stages, iteration)
if not stage.streaming_training:
self._streaming_state = None
return self._single_train_step_short(batch, iteration, stage)
return self._single_train_step_streaming(batch, iteration, stage)
def _single_train_step_short(
self,
batch: dict[str, Any],
iteration: int,
stage: DistillStage,
) -> tuple[
dict[str, torch.Tensor],
dict[str, Any],
dict[str, LogScalar],
]:
batch = self._with_padded_first_frame_latent(
batch,
target_latent_frames=int(self.training_config.data.num_latent_t),
)
training_batch = self.student.prepare_batch(
batch,
generator=self.cuda_generator,
latents_source="zeros",
)
self._truncate_batch_latents(training_batch, stage.num_latent_t)
return self._losses_for_batch(
training_batch,
iteration,
stage=stage,
chunk_mask=None,
)
def _single_train_step_streaming(
self,
batch: dict[str, Any],
iteration: int,
stage: DistillStage,
) -> tuple[
dict[str, torch.Tensor],
dict[str, Any],
dict[str, LogScalar],
]:
state = self._ensure_streaming_state(batch, stage)
update_student = self._should_update_student(iteration)
generator_loss = torch.zeros(
(),
device=state.batch.latents.device,
dtype=state.batch.latents.dtype,
)
student_ctx = None
generator_outputs: dict[str, Any] = {}
pred_x0, chunk_mask, chunk_info = self._generate_streaming_chunk(
state,
with_grad=update_student,
)
if update_student:
student_ctx = (
state.batch.timesteps,
state.batch.attn_metadata_vsa,
)
generator_loss = self._dmd_loss_masked(
pred_x0,
state.batch,
chunk_mask=chunk_mask,
)
generator_outputs["dmd_latent_vis_dict"] = {
"generator_pred_video": pred_x0.detach(),
"streaming_chunk_mask": chunk_mask.detach(),
}
self.student.backward(
generator_loss,
student_ctx,
grad_accum_rounds=self._gradient_accumulation_steps(),
)
generator_loss = generator_loss.detach()
student_ctx = None
fake_score_loss, critic_ctx, critic_outputs = (self._critic_flow_matching_loss_for_x0(
pred_x0.detach(),
state.batch,
chunk_mask=chunk_mask,
))
self._log_train_chunk(
iteration=iteration,
stage=stage,
update_student=update_student,
chunk_info=chunk_info,
)
total_loss = generator_loss + fake_score_loss
outputs: dict[str, Any] = dict(generator_outputs)
outputs.update(critic_outputs)
outputs["_fv_backward"] = {
"update_student": False,
"student_ctx": student_ctx,
"critic_ctx": critic_ctx,
}
metrics: dict[str, LogScalar] = {
"update_student": float(update_student),
"distill_stage_index": float(self._distill_stages.index(stage)),
"streaming_current_length": float(state.current_length),
"streaming_max_length": float(self._stage_max_length(stage)),
}
return (
{
"total_loss": total_loss,
"generator_loss": generator_loss,
"fake_score_loss": fake_score_loss,
},
outputs,
metrics,
)
def _losses_for_batch(
self,
training_batch: Any,
iteration: int,
*,
stage: DistillStage,
chunk_mask: torch.Tensor | None,
) -> tuple[
dict[str, torch.Tensor],
dict[str, Any],
dict[str, LogScalar],
]:
update_student = self._should_update_student(iteration)
generator_loss = torch.zeros(
(),
device=training_batch.latents.device,
dtype=training_batch.latents.dtype,
)
student_ctx = None
if update_student:
generator_pred_x0 = self._student_rollout(
training_batch,
with_grad=True,
)
student_ctx = (
training_batch.timesteps,
training_batch.attn_metadata_vsa,
)
generator_loss = self._dmd_loss_masked(
generator_pred_x0,
training_batch,
chunk_mask=chunk_mask,
)
self.student.backward(
generator_loss,
student_ctx,
grad_accum_rounds=self._gradient_accumulation_steps(),
)
generator_loss = generator_loss.detach()
student_ctx = None
with torch.no_grad():
generator_pred_x0 = self._student_rollout(
training_batch,
with_grad=False,
)
fake_score_loss, critic_ctx, critic_outputs = (self._critic_flow_matching_loss_for_x0(
generator_pred_x0,
training_batch,
chunk_mask=chunk_mask,
))
total_loss = generator_loss + fake_score_loss
outputs: dict[str, Any] = dict(critic_outputs)
outputs["_fv_backward"] = {
"update_student": False,
"student_ctx": student_ctx,
"critic_ctx": critic_ctx,
}
metrics: dict[str, LogScalar] = {
"update_student": float(update_student),
"distill_stage_index": float(self._distill_stages.index(stage)),
"active_num_latent_t": float(stage.num_latent_t),
}
self._log_train_chunk(
iteration=iteration,
stage=stage,
update_student=update_student,
chunk_info=_StreamingChunkInfo(
chunk_start=0,
chunk_end=int(stage.num_latent_t) - 1,
train_start=0,
train_end=int(stage.num_latent_t) - 1,
new_frames=int(stage.num_latent_t),
overlap=0,
current_length=int(stage.num_latent_t),
max_length=int(stage.num_latent_t),
),
)
return (
{
"total_loss": total_loss,
"generator_loss": generator_loss,
"fake_score_loss": fake_score_loss,
},
outputs,
metrics,
)
def _gradient_accumulation_steps(self) -> int:
loop_config = getattr(self.training_config, "loop", None)
return max(1, int(getattr(loop_config, "gradient_accumulation_steps", 1) or 1))
def _student_rollout_streaming(
self,
batch: Any,
*,
with_grad: bool,
) -> torch.Tensor:
# Short-stage rollout should respect the truncated batch length above.
return super()._student_rollout_streaming(batch, with_grad=with_grad)
def _ensure_streaming_state(
self,
raw_batch: dict[str, Any],
stage: DistillStage,
) -> _StreamingState:
state = self._streaming_state
if (state is None or state.stage != stage or not self._can_generate_more(state)):
if state is not None:
self.student.clear_caches(cache_tag=state.cache_tag)
raw_batch = self._with_padded_first_frame_latent(
raw_batch,
target_latent_frames=int(self.training_config.data.num_latent_t),
)
training_batch = self.student.prepare_batch(
raw_batch,
generator=self.cuda_generator,
latents_source="zeros",
)
self._truncate_batch_latents(
training_batch,
self._stage_max_length(stage),
)
cache_tag = f"streaming_long_{id(training_batch)}"
self.student.clear_caches(cache_tag=cache_tag)
state = _StreamingState(
stage=stage,
batch=training_batch,
current_length=0,
previous_latents=None,
cache_tag=cache_tag,
)
self._streaming_state = state
return state
def _with_padded_first_frame_latent(
self,
raw_batch: dict[str, Any],
*,
target_latent_frames: int,
) -> dict[str, Any]:
first_frame_latent = raw_batch.get("first_frame_latent")
if first_frame_latent is None:
return raw_batch
if first_frame_latent.ndim != 5:
raise ValueError("first_frame_latent must have shape [B, C, T, H, W], "
f"got {tuple(first_frame_latent.shape)}")
target_latent_frames = int(target_latent_frames)
if first_frame_latent.shape[2] >= target_latent_frames:
return raw_batch
# Long tuning passes first-frame-only MatrixGame I2V condition as a full-length latent tensor.
pad_frames = target_latent_frames - int(first_frame_latent.shape[2])
pad = torch.zeros(
first_frame_latent.shape[0],
first_frame_latent.shape[1],
pad_frames,
first_frame_latent.shape[3],
first_frame_latent.shape[4],
device=first_frame_latent.device,
dtype=first_frame_latent.dtype,
)
batch = dict(raw_batch)
batch["first_frame_latent"] = torch.cat([first_frame_latent, pad], dim=2)
return batch
def _can_generate_more(self, state: _StreamingState) -> bool:
max_length = self._stage_max_length(state.stage)
return state.current_length < max_length
def _generate_streaming_chunk(
self,
state: _StreamingState,
*,
with_grad: bool,
) -> tuple[torch.Tensor, torch.Tensor, _StreamingChunkInfo]:
assert isinstance(self.student, CausalModelBase)
batch = state.batch
latents = batch.latents
if latents is None:
raise RuntimeError("TrainingBatch.latents is required")
device = latents.device
dtype = latents.dtype
batch_size = int(latents.shape[0])
max_length = self._stage_max_length(state.stage)
chunk_size = self._stage_chunk_size(state.stage)
current = int(state.current_length)
remaining = max_length - current
if remaining <= 0:
raise RuntimeError("Streaming sequence has no remaining frames")
if state.previous_latents is None:
new_frames = min(chunk_size, remaining)
overlap = 0
else:
new_frames = self._select_new_frame_count(
state.stage,
remaining=remaining,
device=device,
)
overlap = min(
chunk_size - new_frames,
int(state.previous_latents.shape[1]),
)
if new_frames <= 0:
raise RuntimeError("Streaming new frame count must be positive")
if current + new_frames > max_length:
raise RuntimeError("Streaming chunk would exceed max length")
rollout_block_size = int(self._chunk_size)
if rollout_block_size <= 0:
raise ValueError("method.chunk_size must be positive")
if self._streaming_require_full_blocks and new_frames % rollout_block_size != 0:
raise ValueError("Streaming new-frame count must be divisible by "
"method.chunk_size: "
f"new_frames={new_frames}, "
f"chunk_size={rollout_block_size}")
denoising_steps = self._get_denoising_step_list(device)
num_steps = int(denoising_steps.numel())
noise_full = torch.randn(
(
batch_size,
new_frames,
int(latents.shape[2]),
int(latents.shape[3]),
int(latents.shape[4]),
),
device=device,
dtype=dtype,
generator=self.cuda_generator,
)
block_lengths: list[int] = []
remaining_block_frames = int(new_frames)
while remaining_block_frames > 0:
current_block = min(rollout_block_size, remaining_block_frames)
block_lengths.append(current_block)
remaining_block_frames -= current_block
exit_indices = self._sample_exit_indices(
num_blocks=len(block_lengths),
num_steps=num_steps,
device=device,
)
block0_exit_idx = int(exit_indices[0])
denoised_from, denoised_to = self._get_self_forcing_timestep_bounds(
denoising_steps,
block0_exit_idx,
)
denoised_blocks: list[torch.Tensor] = []
local_start = 0
for block_idx, block_frames in enumerate(block_lengths):
block_start = current + local_start
noisy_block = noise_full[:, local_start:local_start + block_frames]
exit_idx = int(exit_indices[block_idx])
pred_x0_block: torch.Tensor | None = None
for step_idx, current_timestep in enumerate(denoising_steps):
timestep_block = current_timestep * torch.ones(
(batch_size, block_frames),
device=device,
dtype=torch.float32,
)
exit_flag = step_idx == exit_idx
enable_grad = (bool(with_grad) and bool(self._enable_gradient_in_rollout) and torch.is_grad_enabled())
if not exit_flag:
with torch.no_grad():
pred_noise = self.student.predict_noise_streaming(
noisy_block,
timestep_block,
batch,
conditional=True,
cache_tag=state.cache_tag,
store_kv=False,
cur_start_frame=block_start,
cfg_uncond=self._cfg_uncond,
attn_kind="vsa",
)
if pred_noise is None:
raise RuntimeError("predict_noise_streaming returned None")
pred_x0_block = 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_block,
torch.randn(
pred_x0_block.shape,
device=device,
dtype=pred_x0_block.dtype,
generator=self.cuda_generator,
),
next_timestep * torch.ones(
(batch_size, block_frames),
device=device,
dtype=torch.float32,
),
)
else:
next_timestep_block = next_timestep * torch.ones(
(batch_size, block_frames),
device=device,
dtype=torch.float32,
)
sigma_cur = self._timestep_to_sigma(timestep_block).view(
batch_size,
block_frames,
1,
1,
1,
)
sigma_next = self._timestep_to_sigma(next_timestep_block).view(
batch_size,
block_frames,
1,
1,
1,
)
eps = (noisy_block - (1 - sigma_cur) * pred_x0_block) / sigma_cur.clamp_min(1e-8)
noisy_block = (1 - sigma_next) * pred_x0_block + 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=state.cache_tag,
store_kv=False,
cur_start_frame=block_start,
cfg_uncond=self._cfg_uncond,
attn_kind="vsa",
)
if pred_noise is None:
raise RuntimeError("predict_noise_streaming returned None")
pred_x0_block = 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
if pred_x0_block is None:
raise RuntimeError("Streaming rollout produced no block")
denoised_blocks.append(pred_x0_block)
with torch.no_grad():
context_timestep = torch.ones(
(batch_size, block_frames),
device=device,
dtype=torch.float32,
) * float(self._context_noise)
context_latents = self._sf_add_noise(
pred_x0_block.detach(),
torch.randn(
pred_x0_block.shape,
device=device,
dtype=pred_x0_block.dtype,
generator=self.cuda_generator,
),
context_timestep,
)
_ = self.student.predict_noise_streaming(
context_latents,
context_timestep,
batch,
conditional=True,
cache_tag=state.cache_tag,
store_kv=True,
cur_start_frame=block_start,
cfg_uncond=self._cfg_uncond,
attn_kind="vsa",
)
local_start += block_frames
if not denoised_blocks:
raise RuntimeError("Streaming rollout produced no chunk")
pred_x0_chunk = torch.cat(denoised_blocks, dim=1)
batch.denoised_timestep_from = denoised_from
batch.denoised_timestep_to = denoised_to
if overlap > 0 and state.previous_latents is not None:
overlap_latents = state.previous_latents[:, -overlap:].detach()
if self._streaming_reencode_overlap_anchor:
anchor_source = torch.cat(
[state.previous_latents.detach(), pred_x0_chunk],
dim=1,
)
full_chunk = self._process_first_frame_anchor(
anchor_source,
target_latents=chunk_size,
)
else:
full_chunk = torch.cat([overlap_latents, pred_x0_chunk], dim=1)
chunk_mask = torch.zeros(
(batch_size, chunk_size, 1, 1, 1),
device=device,
dtype=torch.bool,
)
chunk_mask[:, overlap:] = True
else:
full_chunk = pred_x0_chunk
chunk_mask = torch.ones(
(batch_size, new_frames, 1, 1, 1),
device=device,
dtype=torch.bool,
)
state.current_length = current + new_frames
temporal_compression_ratio = int(self.training_config.pipeline_config.vae_config.arch_config.
temporal_compression_ratio # type: ignore[union-attr]
)
chunk_info = _StreamingChunkInfo(
chunk_start=current - overlap,
chunk_end=current + new_frames - 1,
train_start=current,
train_end=current + new_frames - 1,
new_frames=new_frames,
overlap=overlap,
current_length=state.current_length,
max_length=max_length,
)
score_frame_start = max(0, current - overlap)
batch.matrixgame_streaming_chunk_info = {
"frame_start": score_frame_start,
"frame_end": score_frame_start + int(full_chunk.shape[1]),
"new_frame_start": int(overlap),
"new_frames": int(new_frames),
"gradient_mask": chunk_mask,
"condition_image_latents": None,
"condition_image_embeds": None,
"generator_chunk_start": int(current),
"generator_chunk_end": int(current + new_frames),
"generator_action_frame_end": int((current + new_frames - 1) * temporal_compression_ratio + 1),
"window_label": f"window_{score_frame_start}_{score_frame_start + int(full_chunk.shape[1])}",
}
state.previous_latents = full_chunk.detach()[:, -chunk_size:]
if not dist.is_initialized() or dist.get_rank() == 0:
logger.info(
"[StreamingLong] stage=%s start=%s new=%s overlap=%s "
"blocks=%s current_length=%s/%s",
state.stage.name,
current,
new_frames,
overlap,
block_lengths,
state.current_length,
max_length,
)
return full_chunk, chunk_mask, chunk_info
def _get_self_forcing_timestep_bounds(
self,
step_list: torch.Tensor,
exit_idx: int,
) -> tuple[int | None, int | None]:
scheduler_timesteps = self._sf_scheduler.timesteps.to(
device=step_list.device,
dtype=torch.float32,
)
exit_idx = min(int(exit_idx), int(step_list.numel()) - 1)
denoising_timestep = step_list[exit_idx].to(dtype=scheduler_timesteps.dtype)
denoised_timestep_from = int(self.student.num_train_timesteps) - int(
torch.argmin(
(scheduler_timesteps - denoising_timestep).abs(),
dim=0,
).item())
if exit_idx == int(step_list.numel()) - 1:
return denoised_timestep_from, 0
next_denoising_timestep = step_list[exit_idx + 1].to(dtype=scheduler_timesteps.dtype)
denoised_timestep_to = int(self.student.num_train_timesteps) - int(
torch.argmin(
(scheduler_timesteps - next_denoising_timestep).abs(),
dim=0,
).item())
return denoised_timestep_from, denoised_timestep_to
def _process_first_frame_anchor(
self,
frames: torch.Tensor,
*,
target_latents: int,
) -> torch.Tensor:
if frames.shape[1] <= 1:
return frames
process_latents = min(int(target_latents), int(frames.shape[1]))
if process_latents <= 1:
return frames[:, -process_latents:]
vae = getattr(self.student, "vae", None)
if vae is None:
raise RuntimeError("Streaming first-frame anchor requires student.vae")
self._require_vae_encoder(vae)
inject_k = min(int(self._streaming_anchor_inject_k), process_latents)
if inject_k == 1:
with torch.no_grad():
latents_to_decode = frames[:, :-(process_latents - 1), ...]
decode_latents = self._denormalize_wan_latents(
latents_to_decode,
vae=vae,
)
pixels = vae.decode(decode_latents)
last_frame_pixel = pixels[:, :, -1:, :, :].float()
image_latent = self._encode_wan_latent(
vae,
last_frame_pixel,
like=frames,
)
remaining = frames[:, -(process_latents - 1):, ...]
return torch.cat([image_latent, remaining], dim=1)
out_slice = frames[:, -process_latents:, ...]
with torch.no_grad():
first_k = out_slice[:, :inject_k, ...]
decode_latents = self._denormalize_wan_latents(
first_k,
vae=vae,
)
pixels = vae.decode(decode_latents).float()
new_first_k = self._encode_wan_latent(
vae,
pixels,
like=frames,
)
remaining = out_slice[:, inject_k:, ...]
return torch.cat([new_first_k, remaining], dim=1)
def _denormalize_wan_latents(
self,
frames: torch.Tensor,
*,
vae: Any,
) -> torch.Tensor:
latents = frames.permute(0, 2, 1, 3, 4).float()
cfg = getattr(vae, "config", None)
if (cfg is not None and hasattr(cfg, "latents_mean") and hasattr(cfg, "latents_std")):
latents_mean = torch.tensor(
cfg.latents_mean,
device=latents.device,
dtype=latents.dtype,
).view(1, -1, 1, 1, 1)
latents_std = torch.tensor(
cfg.latents_std,
device=latents.device,
dtype=latents.dtype,
).view(1, -1, 1, 1, 1)
return latents * latents_std + latents_mean
if hasattr(vae, "scaling_factor"):
scaling_factor = self._vae_tensor_attr(
vae,
"scaling_factor",
like=latents,
default=1.0,
)
latents = latents / scaling_factor
shift_factor = self._vae_tensor_attr(
vae,
"shift_factor",
like=latents,
default=0.0,
)
return latents + shift_factor
def _normalize_wan_latents(
self,
latents: torch.Tensor,
*,
vae: Any,
) -> torch.Tensor:
cfg = getattr(vae, "config", None)
if (cfg is not None and hasattr(cfg, "latents_mean") and hasattr(cfg, "latents_std")):
latents_mean = torch.tensor(
cfg.latents_mean,
device=latents.device,
dtype=latents.dtype,
).view(1, -1, 1, 1, 1)
latents_std = torch.tensor(
cfg.latents_std,
device=latents.device,
dtype=latents.dtype,
).view(1, -1, 1, 1, 1)
return (latents - latents_mean) / latents_std
shift_factor = self._vae_tensor_attr(
vae,
"shift_factor",
like=latents,
default=0.0,
)
latents = latents - shift_factor
if hasattr(vae, "scaling_factor"):
scaling_factor = self._vae_tensor_attr(
vae,
"scaling_factor",
like=latents,
default=1.0,
)
latents = latents * scaling_factor
return latents
def _encode_wan_latent(
self,
vae: Any,
pixels: torch.Tensor,
*,
like: torch.Tensor,
) -> torch.Tensor:
encoded = vae.encode(pixels)
if isinstance(encoded, torch.Tensor):
raw_latent = encoded
elif hasattr(encoded, "mode"):
raw_latent = encoded.mode()
else:
latent_dist = getattr(encoded, "latent_dist", None)
if latent_dist is not None:
raw_latent = (latent_dist.mode() if hasattr(latent_dist, "mode") else latent_dist.mean)
else:
mean = getattr(encoded, "mean", None)
if not isinstance(mean, torch.Tensor):
raise RuntimeError("Unsupported VAE encode output")
raw_latent = mean
raw_latent = raw_latent.to(device=like.device, dtype=torch.float32)
normalized = self._normalize_wan_latents(
raw_latent,
vae=vae,
)
return normalized.permute(0, 2, 1, 3, 4).to(dtype=like.dtype)
def _require_vae_encoder(self, vae: Any) -> None:
target = getattr(vae, "module", vae)
encoder = getattr(target, "encoder", None)
if encoder is None:
raise RuntimeError("Streaming first-frame anchor requires a VAE "
"loaded with encoder weights")
if (getattr(target, "use_feature_cache", False) and not hasattr(target, "_enc_feat_map")):
config = getattr(target, "config", None)
if config is not None:
config.load_encoder = True
clear_cache = getattr(target, "clear_cache", None)
if callable(clear_cache):
clear_cache()
if (getattr(target, "use_feature_cache", False) and not hasattr(target, "_enc_feat_map")):
raise RuntimeError("Streaming first-frame anchor requires VAE "
"encoder feature-cache buffers")
def _vae_tensor_attr(
self,
vae: Any,
name: str,
*,
like: torch.Tensor,
default: float,
) -> torch.Tensor:
value = getattr(vae, name, None)
if value is None and hasattr(vae, "module"):
value = getattr(vae.module, name, None)
if value is None:
value = default
if isinstance(value, torch.Tensor):
return value.to(device=like.device, dtype=like.dtype)
return torch.tensor(value, device=like.device, dtype=like.dtype)
def _log_train_chunk(
self,
*,
iteration: int,
stage: DistillStage,
update_student: bool,
chunk_info: _StreamingChunkInfo,
) -> None:
if dist.is_initialized() and dist.get_rank() != 0:
return
generator_state = "on" if update_student else "off"
train_target = "generator+critic" if update_student else "critic"
logger.info(
"[StreamingLong] step=%s stage=%s chunk_latents=%s-%s "
"train_latents=%s-%s new=%s overlap=%s current_length=%s/%s "
"train=%s generator=%s critic=on",
iteration,
stage.name,
chunk_info.chunk_start,
chunk_info.chunk_end,
chunk_info.train_start,
chunk_info.train_end,
chunk_info.new_frames,
chunk_info.overlap,
chunk_info.current_length,
chunk_info.max_length,
train_target,
generator_state,
)
def _select_new_frame_count(
self,
stage: DistillStage,
*,
remaining: int,
device: torch.device,
) -> int:
chunk_size = self._stage_chunk_size(stage)
fixed_overlap = stage.streaming_fixed_overlap_latents
if fixed_overlap is None:
fixed_overlap = get_optional_int(
self.method_config,
"streaming_fixed_overlap_latents",
where="method.streaming_fixed_overlap_latents",
)
if fixed_overlap is not None:
return min(chunk_size - int(fixed_overlap), remaining)
min_new = stage.streaming_min_new_frame
if min_new is None:
min_new = get_optional_int(
self.method_config,
"streaming_min_new_frame",
where="method.streaming_min_new_frame",
)
if min_new is None:
min_new = chunk_size
min_new = max(1, int(min_new))
max_new = min(chunk_size, remaining)
if max_new <= min_new:
return max_new
# Match LongLive's step-by-num_frame_per_block selection. With the
# official 21/18 setting this chooses 18 until the final short chunk.
block = int(self._chunk_size)
choices = list(range(min_new, max_new, block))
if not choices:
choices = [max_new]
if not dist.is_initialized() or dist.get_rank() == 0:
idx = torch.randint(
low=0,
high=len(choices),
size=(1, ),
device=device,
dtype=torch.long,
generator=self.cuda_generator,
)
else:
idx = torch.empty((1, ), device=device, dtype=torch.long)
if dist.is_initialized():
dist.broadcast(idx, src=0)
return int(choices[int(idx.item())])
def _stage_max_length(self, stage: DistillStage) -> int:
return int(stage.streaming_max_length or stage.num_latent_t)
def _stage_chunk_size(self, stage: DistillStage) -> int:
chunk_size = stage.streaming_chunk_size
if chunk_size is None:
chunk_size = get_optional_int(
self.method_config,
"streaming_chunk_size",
where="method.streaming_chunk_size",
)
if chunk_size is None:
chunk_size = self._chunk_size
if int(chunk_size) <= 0:
raise ValueError("streaming_chunk_size must be positive")
return int(chunk_size)
def _truncate_batch_latents(self, batch: Any, num_latent_t: int) -> None:
num_latent_t = int(num_latent_t)
if batch.latents is not None:
batch.latents = batch.latents[:, :num_latent_t]
batch.raw_latent_shape = (
int(batch.latents.shape[0]),
int(batch.latents.shape[2]),
int(batch.latents.shape[1]),
int(batch.latents.shape[3]),
int(batch.latents.shape[4]),
)
if batch.noise_latents is not None:
batch.noise_latents = batch.noise_latents[:, :num_latent_t]
if batch.noisy_model_input is not None:
batch.noisy_model_input = batch.noisy_model_input[:, :num_latent_t]
if batch.noise is not None:
batch.noise = batch.noise[:, :num_latent_t]
if batch.timesteps is not None and batch.timesteps.ndim == 2:
batch.timesteps = batch.timesteps[:, :num_latent_t]
if batch.sigmas is not None and batch.sigmas.ndim >= 1:
if batch.sigmas.ndim == 1:
batch.sigmas = batch.sigmas[:num_latent_t]
else:
batch.sigmas = batch.sigmas[:, :num_latent_t]
self._rebuild_attention_metadata(batch)
def _rebuild_attention_metadata(self, batch: Any) -> None:
build_metadata = getattr(self.student, "_build_attention_metadata", None)
if not callable(build_metadata):
return
batch.attn_metadata = None
batch.attn_metadata_vsa = None
build_metadata(batch)
batch.attn_metadata_vsa = copy.copy(batch.attn_metadata)
if batch.attn_metadata is not None:
batch.attn_metadata.VSA_sparsity = 0.0
def _critic_flow_matching_loss_for_x0(
self,
generator_pred_x0: torch.Tensor,
batch: Any,
*,
chunk_mask: torch.Tensor | None,
) -> tuple[torch.Tensor, Any, dict[str, Any]]:
device = generator_pred_x0.device
batch_size = int(generator_pred_x0.shape[0])
fake_score_timestep = self._sample_score_timestep(
batch_size,
device=device,
)
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 = self._masked_mean(
(pred_noise - target)**2,
chunk_mask,
)
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_masked(
self,
generator_pred_x0: torch.Tensor,
batch: Any,
*,
chunk_mask: torch.Tensor | None,
) -> 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
batch_size = int(generator_pred_x0.shape[0])
with torch.no_grad():
timestep = self._sample_score_timestep(
batch_size,
device=device,
)
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",
)
if self._real_score_cfg:
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
else:
real_cfg_x0 = real_cond_x0
denom = torch.abs(generator_pred_x0.float() - real_cfg_x0.float(), ).mean(dim=[1, 2, 3, 4],
keepdim=True).clamp_min(1e-6)
grad = (faker_x0 - real_cfg_x0) / denom
grad = torch.nan_to_num(grad)
loss = 0.5 * (generator_pred_x0.float() - (generator_pred_x0.float() - grad.float()).detach())**2
return self._masked_mean(loss, chunk_mask)
def _score_timestep_from_ratio(self, key: str) -> int | None:
ratio = get_optional_float(
self.method_config,
key,
where=f"method.{key}",
)
if ratio is None:
return None
if not 0.0 <= ratio <= 1.0:
raise ValueError(f"method.{key} must be in [0, 1], got {ratio}")
return int(round(float(ratio) * int(self.student.num_train_timesteps)))
def _sample_score_timestep(
self,
batch_size: int,
*,
device: torch.device,
) -> torch.Tensor:
timestep = torch.randint(
0,
int(self.student.num_train_timesteps),
[batch_size],
device=device,
dtype=torch.long,
generator=self.cuda_generator,
)
timestep = self.student.shift_and_clamp_timestep(timestep)
min_timestep = self._score_min_timestep
max_timestep = self._score_max_timestep
if min_timestep is not None or max_timestep is not None:
timestep = timestep.clamp(
min=0 if min_timestep is None else int(min_timestep),
max=(int(self.student.num_train_timesteps) if max_timestep is None else int(max_timestep)),
)
return timestep
def _masked_mean(
self,
values: torch.Tensor,
mask: torch.Tensor | None,
) -> torch.Tensor:
if mask is None:
return values.mean()
mask_f = mask.to(device=values.device, dtype=values.dtype)
while mask_f.ndim < values.ndim:
mask_f = mask_f.unsqueeze(-1)
denom = mask_f.expand_as(values).sum().clamp_min(1.0)
return (values * mask_f).sum() / denom</code></pre></details>
fastvideo.train.methods.distribution_matching.streaming_long_tuning.StreamingLongTuningMethod.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],
]:
stage = select_distill_stage(self._distill_stages, iteration)
if not stage.streaming_training:
self._streaming_state = None
return self._single_train_step_short(batch, iteration, stage)
return self._single_train_step_streaming(batch, iteration, stage)