From 9edcc0adb4f01f2bc1d040f6d21cd4793268d997 Mon Sep 17 00:00:00 2001 From: =?UTF-8?q?=E2=80=9CBrianChen1129=E2=80=9D?= Date: Sat, 19 Jul 2025 04:31:19 +0000 Subject: [PATCH] update --- .../scheduling_flow_match_euler_discrete.py | 298 ------------------ .../v1/training/distillation_pipeline.py | 14 +- .../v1/training/wan_distillation_pipeline.py | 6 +- 3 files changed, 8 insertions(+), 310 deletions(-) diff --git a/fastvideo/v1/models/schedulers/scheduling_flow_match_euler_discrete.py b/fastvideo/v1/models/schedulers/scheduling_flow_match_euler_discrete.py index 1f50890b3..0eb7cb7a9 100644 --- a/fastvideo/v1/models/schedulers/scheduling_flow_match_euler_discrete.py +++ b/fastvideo/v1/models/schedulers/scheduling_flow_match_euler_discrete.py @@ -587,301 +587,3 @@ class FlowMatchEulerDiscreteScheduler(SchedulerMixin, ConfigMixin, BaseScheduler def __len__(self): return self.config.num_train_timesteps - -class FlowMatchScheduler(): - order = 1 - def __init__(self, num_inference_steps=100, num_train_timesteps=1000, shift=3.0, sigma_max=1.0, sigma_min=0.003 / 1.002, inverse_timesteps=False, extra_one_step=False, reverse_sigmas=False): - self.num_train_timesteps = num_train_timesteps - self.shift = shift - self.sigma_max = sigma_max - self.sigma_min = sigma_min - self.inverse_timesteps = inverse_timesteps - self.extra_one_step = extra_one_step - self.reverse_sigmas = reverse_sigmas - self.set_timesteps(num_inference_steps) - - def set_timesteps(self, num_inference_steps=100, denoising_strength=1.0, training=False): - sigma_start = self.sigma_min + \ - (self.sigma_max - self.sigma_min) * denoising_strength - if self.extra_one_step: - self.sigmas = torch.linspace( - sigma_start, self.sigma_min, num_inference_steps + 1)[:-1] - else: - self.sigmas = torch.linspace( - sigma_start, self.sigma_min, num_inference_steps) - if self.inverse_timesteps: - self.sigmas = torch.flip(self.sigmas, dims=[0]) - self.sigmas = self.shift * self.sigmas / \ - (1 + (self.shift - 1) * self.sigmas) - if self.reverse_sigmas: - self.sigmas = 1 - self.sigmas - self.timesteps = self.sigmas * self.num_train_timesteps - if training: - x = self.timesteps - y = torch.exp(-2 * ((x - num_inference_steps / 2) / - num_inference_steps) ** 2) - y_shifted = y - y.min() - bsmntw_weighing = y_shifted * \ - (num_inference_steps / y_shifted.sum()) - self.linear_timesteps_weights = bsmntw_weighing - - def step(self, model_output, timestep, sample, to_final=False): - self.sigmas = self.sigmas.to(model_output.device) - self.timesteps = self.timesteps.to(model_output.device) - timestep_id = torch.argmin( - (self.timesteps.unsqueeze(0) - timestep.unsqueeze(1)).abs(), dim=1) - sigma = self.sigmas[timestep_id].reshape(-1, 1, 1, 1) - if to_final or (timestep_id + 1 >= len(self.timesteps)).any(): - sigma_ = 1 if ( - self.inverse_timesteps or self.reverse_sigmas) else 0 - else: - sigma_ = self.sigmas[timestep_id + 1].reshape(-1, 1, 1, 1) - prev_sample = sample + model_output * (sigma_ - sigma) - return prev_sample - - def add_noise(self, original_samples, noise, timestep): - """ - Diffusion forward corruption process. - Input: - - clean_latent: the clean latent with shape [B, C, H, W] - - noise: the noise with shape [B, C, H, W] - - timestep: the timestep with shape [B] - Output: the corrupted latent with shape [B, C, H, W] - """ - self.sigmas = self.sigmas.to(noise.device) - self.timesteps = self.timesteps.to(noise.device) - timestep_id = torch.argmin( - (self.timesteps.unsqueeze(0) - timestep.unsqueeze(1)).abs(), dim=1) - sigma = self.sigmas[timestep_id].reshape(-1, 1, 1, 1) # [21, 1, 1, 1] - sample = (1 - sigma) * original_samples + sigma * noise - return sample.type_as(noise) - - def training_target(self, sample, noise, timestep): - target = noise - sample - return target - - def training_weight(self, timestep): - timestep_id = torch.argmin( - (self.timesteps - timestep.to(self.timesteps.device)).abs()) - weights = self.linear_timesteps_weights[timestep_id] - return weights - -# class FlowMatchDiscreteScheduler(SchedulerMixin, ConfigMixin, BaseScheduler): -# """ -# Euler scheduler. - -# This model inherits from [`SchedulerMixin`] and [`ConfigMixin`]. Check the superclass documentation for the generic -# methods the library implements for all schedulers such as loading and saving. - -# Args: -# num_train_timesteps (`int`, defaults to 1000): -# The number of diffusion steps to train the model. -# timestep_spacing (`str`, defaults to `"linspace"`): -# The way the timesteps should be scaled. Refer to Table 2 of the [Common Diffusion Noise Schedules and -# Sample Steps are Flawed](https://huggingface.co/papers/2305.08891) for more information. -# shift (`float`, defaults to 1.0): -# The shift value for the timestep schedule. -# reverse (`bool`, defaults to `True`): -# Whether to reverse the timestep schedule. -# """ - -# _compatibles: list[Any] = [] -# order = 1 - -# @register_to_config -# def __init__( -# self, -# num_train_timesteps: int = 1000, -# shift: float = 1.0, -# reverse: bool = True, -# solver: str = "euler", -# n_tokens: Optional[int] = None, -# **kwargs, -# ): -# sigmas = torch.linspace(1, 0, num_train_timesteps + 1) - -# if not reverse: -# sigmas = sigmas.flip(0) - -# self.sigmas = sigmas -# # the value fed to model -# self.timesteps = (sigmas[:-1] * -# num_train_timesteps).to(dtype=torch.float32) - -# self._step_index: int | None = None -# self._begin_index = 0 - -# self.supported_solver = ["euler"] -# if solver not in self.supported_solver: -# raise ValueError( -# f"Solver {solver} not supported. Supported solvers: {self.supported_solver}" -# ) - -# BaseScheduler.__init__(self) - -# @property -# def step_index(self): -# """ -# The index counter for current timestep. It will increase 1 after each scheduler step. -# """ -# return self._step_index - -# @property -# def begin_index(self): -# """ -# The index for the first timestep. It should be set from pipeline with `set_begin_index` method. -# """ -# return self._begin_index - -# # Copied from diffusers.schedulers.scheduling_dpmsolver_multistep.DPMSolverMultistepScheduler.set_begin_index -# def set_begin_index(self, begin_index: int = 0): -# """ -# Sets the begin index for the scheduler. This function should be run from pipeline before the inference. - -# Args: -# begin_index (`int`): -# The begin index for the scheduler. -# """ -# self._begin_index = begin_index - -# def _sigma_to_t(self, sigma): -# return sigma * self.config.num_train_timesteps - -# def set_timesteps( -# self, -# num_inference_steps: int, -# device: Union[str, torch.device] = None, -# n_tokens: int = 0, -# ): -# """ -# Sets the discrete timesteps used for the diffusion chain (to be run before inference). - -# Args: -# num_inference_steps (`int`): -# The number of diffusion steps used when generating samples with a pre-trained model. -# device (`str` or `torch.device`, *optional*): -# The device to which the timesteps should be moved to. If `None`, the timesteps are not moved. -# n_tokens (`int`, *optional*): -# Number of tokens in the input sequence. -# """ -# self.num_inference_steps = num_inference_steps - -# sigmas = torch.linspace(1, 0, num_inference_steps + 1) -# sigmas = self.sd3_time_shift(sigmas) - -# if not self.config.reverse: -# sigmas = 1 - sigmas - -# self.sigmas = sigmas -# if not getattr(self.config, "timesteps_scale", True): -# self.timesteps = sigmas[:-1] # for stepvideo -# else: -# self.timesteps = (sigmas[:-1] * self.config.num_train_timesteps).to( -# dtype=torch.float32, device=device) -# # Reset step index -# self._step_index = None - -# def index_for_timestep(self, timestep, schedule_timesteps=None) -> int: -# if schedule_timesteps is None: -# schedule_timesteps = self.timesteps - -# indices = (schedule_timesteps == timestep).nonzero() - -# # The sigma index that is taken for the **very** first `step` -# # is always the second index (or the last index if there is only 1) -# # This way we can ensure we don't accidentally skip a sigma in -# # case we start in the middle of the denoising schedule (e.g. for image-to-image) -# pos = 1 if len(indices) > 1 else 0 - -# idx: int = indices[pos].item() - -# return idx - -# def set_shift(self, shift: float) -> None: -# self.config.shift = shift - -# def set_timesteps_scale(self, timesteps_scale: bool) -> None: -# self.config.timesteps_scale = timesteps_scale - -# def _init_step_index(self, timestep) -> None: -# if self.begin_index is None: -# if isinstance(timestep, torch.Tensor): -# timestep = timestep.to(self.timesteps.device) -# self._step_index = self.index_for_timestep(timestep) -# else: -# self._step_index = self._begin_index - -# def scale_model_input(self, -# sample: torch.Tensor, -# timestep: Optional[int] = None) -> torch.Tensor: -# return sample - -# def sd3_time_shift(self, t: torch.Tensor): -# return (self.config.shift * t) / (1 + (self.config.shift - 1) * t) - -# def step( -# self, -# model_output: torch.FloatTensor, -# timestep: Union[float, torch.FloatTensor], -# sample: torch.FloatTensor, -# return_dict: bool = True, -# **kwargs, -# ) -> Union[FlowMatchDiscreteSchedulerOutput, Tuple]: -# """ -# Predict the sample from the previous timestep by reversing the SDE. This function propagates the diffusion -# process from the learned model outputs (most often the predicted noise). - -# Args: -# model_output (`torch.FloatTensor`): -# The direct output from learned diffusion model. -# timestep (`float`): -# The current discrete timestep in the diffusion chain. -# sample (`torch.FloatTensor`): -# A current instance of a sample created by the diffusion process. -# generator (`torch.Generator`, *optional*): -# A random number generator. -# n_tokens (`int`, *optional*): -# Number of tokens in the input sequence. -# return_dict (`bool`): -# Whether or not to return a [`~schedulers.scheduling_euler_discrete.EulerDiscreteSchedulerOutput`] or -# tuple. - -# Returns: -# [`~schedulers.scheduling_euler_discrete.EulerDiscreteSchedulerOutput`] or `tuple`: -# If return_dict is `True`, [`~schedulers.scheduling_euler_discrete.EulerDiscreteSchedulerOutput`] is -# returned, otherwise a tuple is returned where the first element is the sample tensor. -# """ - -# if isinstance(timestep, (int, torch.IntTensor, torch.LongTensor)): -# raise ValueError(( -# "Passing integer indices (e.g. from `enumerate(timesteps)`) as timesteps to" -# " `EulerDiscreteScheduler.step()` is not supported. Make sure to pass" -# " one of the `scheduler.timesteps` as a timestep."), ) - -# if self.step_index is None: -# self._init_step_index(timestep) - -# # Upcast to avoid precision issues when computing prev_sample -# sample = sample.to(torch.float32) - -# assert self.step_index is not None -# dt = self.sigmas[self.step_index + 1] - self.sigmas[self.step_index] - -# if self.config.solver == "euler": -# prev_sample = sample + model_output.to(torch.float32) * dt -# else: -# raise ValueError( -# f"Solver {self.config.solver} not supported. Supported solvers: {self.supported_solver}" -# ) - -# # upon completion increase step index by one -# assert self._step_index is not None -# self._step_index += 1 - -# if not return_dict: -# return (prev_sample, ) - -# return FlowMatchDiscreteSchedulerOutput(prev_sample=prev_sample) - -# def __len__(self): -# return self.config.num_train_timesteps diff --git a/fastvideo/v1/training/distillation_pipeline.py b/fastvideo/v1/training/distillation_pipeline.py index 8b0d56437..2f4964aa9 100644 --- a/fastvideo/v1/training/distillation_pipeline.py +++ b/fastvideo/v1/training/distillation_pipeline.py @@ -34,7 +34,7 @@ from fastvideo.v1.training.training_utils import ( compute_density_for_timestep_sampling, get_sigmas, load_checkpoint, normalize_dit_input, save_checkpoint, shard_latents_across_sp, prepare_for_saving) from fastvideo.v1.utils import set_random_seed, is_vsa_available -from fastvideo.v1.models.schedulers.scheduling_flow_match_euler_discrete import FlowMatchEulerDiscreteScheduler, FlowMatchScheduler +from fastvideo.v1.models.schedulers.scheduling_flow_match_euler_discrete import FlowMatchEulerDiscreteScheduler from fastvideo.v1.training.activation_checkpoint import ( apply_activation_checkpointing) from fastvideo.v1.attention.backends.video_sparse_attn import ( @@ -76,11 +76,7 @@ class DistillationPipeline(TrainingPipeline): self.vae.requires_grad_(False) self.timestep_shift = self.training_args.pipeline_config.flow_shift - # self.noise_scheduler = FlowMatchEulerDiscreteScheduler(shift=self.timestep_shift) - self.noise_scheduler = FlowMatchScheduler( - shift=8.0, sigma_min=0.0, extra_one_step=True - ) - self.noise_scheduler.set_timesteps(1000, training=True) + self.noise_scheduler = FlowMatchEulerDiscreteScheduler(shift=self.timestep_shift) # 2. Distillation-specific initialization # The parent class already sets self.transformer as the student model @@ -454,7 +450,7 @@ class DistillationPipeline(TrainingPipeline): assert self.training_args is not None training_batch = self._prepare_distillation(training_batch) - TRAIN_STUDENT = self.current_trainstep % self.student_critic_update_ratio == 0 + TRAIN_STUDENT = (self.current_trainstep % self.student_critic_update_ratio == 0) or (self.current_trainstep == 1) # for _ in range(self.training_args.gradient_accumulation_steps): training_batch = self._get_next_batch(training_batch) @@ -797,7 +793,7 @@ class DistillationPipeline(TrainingPipeline): step_times: deque[float] = deque(maxlen=100) self._log_training_info() - self._log_validation(self.student_transformer, self.training_args, -1) + self._log_validation(self.student_transformer, self.training_args, 0) progress_bar = tqdm( range(0, self.training_args.max_train_steps), @@ -806,7 +802,7 @@ class DistillationPipeline(TrainingPipeline): disable=self.local_rank > 0, ) - for step in range(self.init_steps, + for step in range(self.init_steps + 1, self.training_args.max_train_steps + 1): start_time = time.perf_counter() current_vsa_sparsity = self.training_args.VSA_sparsity if vsa_available else 0.0 diff --git a/fastvideo/v1/training/wan_distillation_pipeline.py b/fastvideo/v1/training/wan_distillation_pipeline.py index fbae6cdba..7d83d76a4 100644 --- a/fastvideo/v1/training/wan_distillation_pipeline.py +++ b/fastvideo/v1/training/wan_distillation_pipeline.py @@ -6,8 +6,8 @@ import torch from fastvideo.v1.distributed import get_local_torch_device from fastvideo.v1.fastvideo_args import FastVideoArgs, TrainingArgs from fastvideo.v1.logger import init_logger -from fastvideo.v1.models.schedulers.scheduling_flow_unipc_multistep import ( - FlowUniPCMultistepScheduler) +from fastvideo.v1.models.schedulers.scheduling_flow_match_euler_discrete import ( + FlowMatchEulerDiscreteScheduler) from fastvideo.v1.pipelines.wan.wan_dmd_pipeline import WanDmdPipeline from fastvideo.v1.training.distillation_pipeline import DistillationPipeline from fastvideo.v1.pipelines.pipeline_batch_info import (ForwardBatch, @@ -29,7 +29,7 @@ class WanDistillationPipeline(DistillationPipeline): def initialize_pipeline(self, fastvideo_args: FastVideoArgs): """Initialize Wan-specific scheduler.""" - self.modules["scheduler"] = FlowUniPCMultistepScheduler( + self.modules["scheduler"] = FlowMatchEulerDiscreteScheduler( shift=fastvideo_args.pipeline_config.flow_shift) def create_training_stages(self, training_args: TrainingArgs):