--- title: 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) -> _FastvideoLogger

The 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.Tensor

Convert 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.Tensor

the 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') -> None

Clear 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 | None

Streaming 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 | None

Predict 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 | None

fastvideo.train.methods.distribution_matching.streaming_long_tuning.get_optional_int

get_optional_int(mapping: dict[str, Any], key: str, *, where: str) -> int | None

fastvideo.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

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

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]

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

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)