update df scheduler
This commit is contained in:
H1yori233
@@ -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()
|
||||
})
|
||||
})
|
||||
@@ -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
|
||||
@@ -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
|
||||
|
||||
Reference in New Issue
Block a user