From 2ace7dc6f4da92364f8f0c8fdf081da23907b9e6 Mon Sep 17 00:00:00 2001 From: H1yori233 Date: Mon, 16 Mar 2026 15:49:09 -0700 Subject: [PATCH] update df scheduler --- fastvideo/models/registry.py | 5 +- .../scheduling_diffusion_forcing.py | 205 ++++++++++++++++++ fastvideo/train/methods/fine_tuning/dfsft.py | 107 ++++++--- 3 files changed, 283 insertions(+), 34 deletions(-) create mode 100644 fastvideo/models/schedulers/scheduling_diffusion_forcing.py diff --git a/fastvideo/models/registry.py b/fastvideo/models/registry.py index f2de89eff..46e8efd1d 100644 --- a/fastvideo/models/registry.py +++ b/fastvideo/models/registry.py @@ -99,6 +99,9 @@ _SCHEDULERS = { "FlowMatchEulerDiscreteScheduler": ("schedulers", "scheduling_flow_match_euler_discrete", "FlowMatchEulerDiscreteScheduler"), + "DiffusionForcingScheduler": + ("schedulers", "scheduling_diffusion_forcing", + "DiffusionForcingScheduler"), "UniPCMultistepScheduler": ("schedulers", "scheduling_unipc_multistep", "UniPCMultistepScheduler"), "FlowUniPCMultistepScheduler": @@ -457,4 +460,4 @@ ModelRegistry = _ModelRegistry({ ) for model_arch, (component_name, mod_relname, cls_name) in _FAST_VIDEO_MODELS.items() -}) +}) \ No newline at end of file diff --git a/fastvideo/models/schedulers/scheduling_diffusion_forcing.py b/fastvideo/models/schedulers/scheduling_diffusion_forcing.py new file mode 100644 index 000000000..0fe15bc23 --- /dev/null +++ b/fastvideo/models/schedulers/scheduling_diffusion_forcing.py @@ -0,0 +1,205 @@ +# SPDX-License-Identifier: Apache-2.0 +from diffusers.configuration_utils import ConfigMixin, register_to_config +from diffusers.schedulers.scheduling_utils import SchedulerMixin +from diffusers.utils import BaseOutput +import torch + +from fastvideo.models.schedulers.base import BaseScheduler + + +class DiffusionForcingSchedulerOutput(BaseOutput): + prev_sample: torch.FloatTensor + + +class DiffusionForcingScheduler(BaseScheduler, ConfigMixin, SchedulerMixin): + + config_name = "scheduler_config.json" + order = 1 + + @register_to_config + def __init__( + self, + num_inference_steps: int = 100, + num_train_timesteps: int = 1000, + shift: float = 5.0, + sigma_max: float = 1.0, + sigma_min: float = 0.0, + extra_one_step: bool = True, + training: bool = False, + ): + self.num_train_timesteps = num_train_timesteps + self.shift = shift + self.sigma_max = sigma_max + self.sigma_min = sigma_min + self.extra_one_step = extra_one_step + self.set_timesteps(num_inference_steps, training=training) + + def sigma_from_timestep(self, timestep: torch.Tensor) -> torch.Tensor: + if not torch.is_tensor(timestep): + timestep = torch.as_tensor(timestep, dtype=torch.float32) + timestep = self._flatten_timestep(timestep) + device = timestep.device + self.sigmas = self.sigmas.to(device) + timestep_id = self._lookup_timestep_indices( + timestep=timestep, + device=device, + ) + return self.sigmas[timestep_id] + + def _flatten_timestep(self, timestep: torch.Tensor) -> torch.Tensor: + if timestep.ndim == 2: + timestep = timestep.flatten(0, 1) + elif timestep.ndim == 0: + timestep = timestep.unsqueeze(0) + elif timestep.ndim != 1: + raise ValueError("timestep must be scalar, [B], [B, T], " + "or [B*T]") + return timestep.to(torch.float32) + + def _lookup_timestep_indices( + self, + *, + timestep: torch.Tensor, + device: torch.device, + ) -> torch.Tensor: + self.timesteps = self.timesteps.to(device) + return torch.argmin( + (self.timesteps.unsqueeze(0) - timestep.unsqueeze(1)).abs(), + dim=1, + ) + + def set_timesteps( + self, + num_inference_steps: int = 100, + denoising_strength: float = 1.0, + training: bool = False, + return_dict: bool = False, + **kwargs, + ): + sigma_start = self.sigma_min + ( + self.sigma_max - self.sigma_min + ) * denoising_strength + if self.extra_one_step: + sigmas = torch.linspace( + sigma_start, self.sigma_min, num_inference_steps + 1 + )[:-1] + else: + sigmas = torch.linspace( + sigma_start, self.sigma_min, num_inference_steps + ) + if self.shift != 1.0: + sigmas = self.shift * sigmas / (1 + (self.shift - 1) * sigmas) + self.sigmas = sigmas.to(torch.float32) + self.timesteps = ( + self.sigmas * float(self.num_train_timesteps) + ).to(torch.float32) + if training: + x = self.timesteps + y = torch.exp( + -2 * ((x - num_inference_steps / 2) / num_inference_steps) ** 2 + ) + y_shifted = y - y.min() + self.linear_timesteps_weights = ( + y_shifted * (num_inference_steps / y_shifted.sum()) + ) + + def step( + self, + model_output: torch.FloatTensor, + timestep: torch.FloatTensor, + sample: torch.FloatTensor, + to_final: bool = False, + return_dict: bool = False, + **kwargs, + ): + timestep = self._flatten_timestep(timestep).to( + model_output.device, + dtype=torch.float32, + ) + self.sigmas = self.sigmas.to(model_output.device) + timestep_id = self._lookup_timestep_indices( + timestep=timestep, + device=model_output.device, + ) + sigma = self.sigmas[timestep_id].reshape(-1, 1, 1, 1) + if to_final: + sigma_next = torch.zeros_like(sigma) + else: + sigma_next = torch.zeros_like(sigma) + valid = timestep_id + 1 < len(self.timesteps) + if valid.any(): + sigma_next[valid] = self.sigmas[ + timestep_id[valid] + 1 + ].reshape(-1, 1, 1, 1) + + prev_sample = sample + model_output * (sigma_next - sigma) + if isinstance(prev_sample, (torch.Tensor, float)) and not return_dict: + return (prev_sample,) + return DiffusionForcingSchedulerOutput(prev_sample=prev_sample) + + @staticmethod + def calculate_alpha_beta_high(sigma, sigma_bound): + alpha = (1 - sigma) / (1 - sigma_bound) + beta = torch.sqrt(sigma**2 - (alpha * sigma_bound) ** 2) + return alpha, beta + + def add_noise(self, original_samples, noise, timestep): + timestep = self._flatten_timestep(timestep).to(noise.device) + self.sigmas = self.sigmas.to(noise.device) + timestep_id = self._lookup_timestep_indices( + timestep=timestep, + device=noise.device, + ) + sigma = self.sigmas[timestep_id].reshape( + -1, 1, 1, 1, + ) + sample = (1 - sigma) * original_samples + sigma * noise + return sample.type_as(noise) + + def add_noise_high( + self, original_samples, noise, timestep, boundary_timestep + ): + timestep = self._flatten_timestep(timestep).to(noise.device) + boundary_timestep = self._flatten_timestep(boundary_timestep).to( + noise.device + ) + self.sigmas = self.sigmas.to(noise.device) + timestep_id = self._lookup_timestep_indices( + timestep=timestep, + device=noise.device, + ) + boundary_timestep_id = self._lookup_timestep_indices( + timestep=boundary_timestep, + device=noise.device, + ) + sigma = self.sigmas[timestep_id].reshape(-1, 1, 1, 1) + sigma_boundary = self.sigmas[boundary_timestep_id].reshape( + -1, 1, 1, 1 + ) + alpha, beta = self.calculate_alpha_beta_high(sigma, sigma_boundary) + sample = alpha * original_samples + beta * noise + return sample.type_as(noise) + + def training_target(self, sample, noise, timestep): + return noise - sample + + def training_weight(self, timestep): + timestep = self._flatten_timestep(timestep) + device = timestep.device + self.linear_timesteps_weights = self.linear_timesteps_weights.to(device) + timestep_id = self._lookup_timestep_indices( + timestep=timestep, + device=device, + ) + return self.linear_timesteps_weights[timestep_id] + + def scale_model_input( + self, sample: torch.Tensor, timestep: int | None = None + ) -> torch.Tensor: + return sample + + def set_shift(self, shift: float) -> None: + self.shift = shift + + +EntryClass = DiffusionForcingScheduler \ No newline at end of file diff --git a/fastvideo/train/methods/fine_tuning/dfsft.py b/fastvideo/train/methods/fine_tuning/dfsft.py index e4eda24d4..3d5f52a8c 100644 --- a/fastvideo/train/methods/fine_tuning/dfsft.py +++ b/fastvideo/train/methods/fine_tuning/dfsft.py @@ -3,11 +3,14 @@ from __future__ import annotations +import copy from typing import Any, Literal import torch import torch.nn.functional as F +from fastvideo.models.schedulers.scheduling_diffusion_forcing import ( + DiffusionForcingScheduler, ) from fastvideo.train.methods.base import TrainingMethod, LogScalar from fastvideo.train.models.base import ModelBase from fastvideo.train.utils.optimizer import ( @@ -31,9 +34,15 @@ class DiffusionForcingSFTMethod(TrainingMethod): raise ValueError("DFSFT requires role 'student'") if not self.student._trainable: raise ValueError("DFSFT requires student to be trainable") + if self.training_config.model.precondition_outputs: + raise ValueError( + "DFSFT only supports official diffusion-forcing loss; " + "set training.model.precondition_outputs=false" + ) self._attn_kind: Literal["dense", "vsa"] = (self._infer_attn_kind()) self._chunk_size = self._parse_chunk_size(self.method_config.get("chunk_size", None)) + self._dfsft_scheduler = self._build_dfsft_scheduler() self._timestep_index_range = (self._parse_timestep_index_range()) # Initialize preprocessors on student. @@ -103,36 +112,30 @@ class DiffusionForcingSFTMethod(TrainingMethod): if (sp_size > 1 and sp_group is not None and hasattr(sp_group, "broadcast")): sp_group.broadcast(timestep_indices, src=0) - scheduler = self.student.noise_scheduler - if scheduler is None: - raise ValueError("DFSFT requires student.noise_scheduler") - - schedule_timesteps = scheduler.timesteps.to(device=clean_latents.device, dtype=torch.float32) - schedule_sigmas = scheduler.sigmas.to( + scheduler = self._dfsft_scheduler + t_inhom = scheduler.timesteps.to( device=clean_latents.device, - dtype=clean_latents.dtype, - ) - t_inhom = schedule_timesteps[timestep_indices] + dtype=torch.float32, + )[timestep_indices] # Override the homogeneous timesteps from prepare_batch # so that set_forward_context (in predict_noise and # backward) receives the correct per-chunk timesteps. training_batch.timesteps = t_inhom + training_batch = self._refresh_attention_metadata(training_batch) - noise = getattr(training_batch, "noise", None) - if noise is None: - noise = torch.randn_like(clean_latents) - else: - if not torch.is_tensor(noise): - raise TypeError("TrainingBatch.noise must be a " - "torch.Tensor when set") - noise = noise.permute(0, 2, 1, 3, 4).to(dtype=clean_latents.dtype) - - noisy_latents = self.student.add_noise( - clean_latents, - noise, - t_inhom.flatten(), + noise = torch.randn( + clean_latents.shape, + generator=self.cuda_generator, + device=clean_latents.device, + dtype=clean_latents.dtype, ) + noisy_latents = scheduler.add_noise( + clean_latents.flatten(0, 1), + noise.flatten(0, 1), + t_inhom.flatten(), + ).unflatten(0, clean_latents.shape[:2]) + training_batch.noise = noise pred = self.student.predict_noise( noisy_latents, @@ -142,14 +145,17 @@ class DiffusionForcingSFTMethod(TrainingMethod): attn_kind=self._attn_kind, ) - if bool(self.training_config.model.precondition_outputs): - sigmas = schedule_sigmas[timestep_indices] - sigmas = sigmas.unsqueeze(-1).unsqueeze(-1).unsqueeze(-1) - pred_x0 = noisy_latents - pred * sigmas - loss = F.mse_loss(pred_x0.float(), clean_latents.float()) - else: - target = noise - clean_latents - loss = F.mse_loss(pred.float(), target.float()) + target = scheduler.training_target(clean_latents, noise, t_inhom) + per_frame_loss = F.mse_loss( + pred.float(), + target.float(), + reduction="none", + ).mean(dim=(2, 3, 4)) + weight = scheduler.training_weight(t_inhom).reshape( + batch_size, + num_latents, + ) + loss = (per_frame_loss * weight.float()).mean() if self._attn_kind == "vsa": attn_metadata = training_batch.attn_metadata_vsa @@ -245,9 +251,7 @@ class DiffusionForcingSFTMethod(TrainingMethod): f"got {type(raw).__name__}") def _parse_timestep_index_range(self, ) -> tuple[int, int]: - scheduler = self.student.noise_scheduler - if scheduler is None: - raise ValueError("DFSFT requires student.noise_scheduler") + scheduler = self._dfsft_scheduler num_steps = int(getattr(scheduler, "config", scheduler).num_train_timesteps) min_ratio = self._parse_ratio( @@ -320,3 +324,40 @@ class DiffusionForcingSFTMethod(TrainingMethod): ) expanded = chunk_indices.repeat_interleave(chunk_size, dim=1) return expanded[:, :num_latents] + + def _build_dfsft_scheduler(self) -> DiffusionForcingScheduler: + student_scheduler = getattr(self.student, "noise_scheduler", None) + if student_scheduler is None: + raise ValueError("DFSFT requires student.noise_scheduler") + num_steps = int( + getattr(student_scheduler, "config", student_scheduler).num_train_timesteps + ) + pipeline_config = self.training_config.pipeline_config + if pipeline_config is None: + raise ValueError("DFSFT requires training_config.pipeline_config") + shift = float( + getattr( + pipeline_config, + "flow_shift", + getattr(self.student, "timestep_shift", 1.0), + ) + ) + scheduler = DiffusionForcingScheduler( + num_inference_steps=num_steps, + num_train_timesteps=num_steps, + shift=shift, + sigma_min=0.0, + extra_one_step=True, + training=True, + ) + scheduler.set_timesteps(num_steps, training=True) + return scheduler + + def _refresh_attention_metadata(self, training_batch: Any) -> Any: + training_batch = self.student._build_attention_metadata(training_batch) + training_batch.attn_metadata_vsa = copy.deepcopy( + training_batch.attn_metadata + ) + if training_batch.attn_metadata is not None: + training_batch.attn_metadata.VSA_sparsity = 0.0 + return training_batch