1472 lines
55 KiB
Python
1472 lines
55 KiB
Python
# SPDX-License-Identifier: Apache-2.0
|
|
"""LongLive-style multi-stage self-forcing for streaming rollouts."""
|
|
|
|
from __future__ import annotations
|
|
|
|
import copy
|
|
from dataclasses import dataclass
|
|
from typing import Any
|
|
|
|
import torch
|
|
import torch.distributed as dist
|
|
|
|
from fastvideo.logger import init_logger
|
|
from fastvideo.models.utils import pred_noise_to_pred_video
|
|
from fastvideo.train.methods.base import LogScalar
|
|
from fastvideo.train.methods.distribution_matching.self_forcing import (
|
|
SelfForcingMethod, )
|
|
from fastvideo.train.models.base import CausalModelBase
|
|
from fastvideo.train.utils.config import get_optional_float, get_optional_int
|
|
|
|
logger = init_logger(__name__)
|
|
|
|
|
|
@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
|
|
|
|
|
|
@dataclass(slots=True)
|
|
class _StreamingState:
|
|
stage: DistillStage
|
|
batch: Any
|
|
current_length: int = 0
|
|
previous_latents: torch.Tensor | None = None
|
|
cache_tag: str = "streaming_long"
|
|
|
|
|
|
@dataclass(frozen=True, slots=True)
|
|
class _StreamingChunkInfo:
|
|
chunk_start: int
|
|
chunk_end: int
|
|
train_start: int
|
|
train_end: int
|
|
new_frames: int
|
|
overlap: int
|
|
current_length: int
|
|
max_length: int
|
|
|
|
|
|
def _as_bool(raw: Any, *, where: str) -> bool:
|
|
if isinstance(raw, bool):
|
|
return raw
|
|
raise ValueError(f"Expected bool at {where}, got {type(raw).__name__}")
|
|
|
|
|
|
def _as_int(raw: Any, *, where: str) -> int:
|
|
if isinstance(raw, bool):
|
|
raise ValueError(f"Expected int at {where}, got bool")
|
|
if isinstance(raw, int):
|
|
return int(raw)
|
|
if isinstance(raw, float) and raw.is_integer():
|
|
return int(raw)
|
|
if isinstance(raw, str) and raw.strip():
|
|
return int(raw)
|
|
raise ValueError(f"Expected int at {where}, got {type(raw).__name__}")
|
|
|
|
|
|
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
|
|
|
|
|
|
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]
|
|
|
|
|
|
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
|