update
This commit is contained in:
“BrianChen1129”
@@ -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
|
||||
|
||||
@@ -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
|
||||
|
||||
@@ -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):
|
||||
|
||||
Reference in New Issue
Block a user