wehub-resource-sync
495 lines
22 KiB
Python
495 lines
22 KiB
Python
# Adopted from https://github.com/guandeh17/Self-Forcing
|
|
# SPDX-License-Identifier: Apache-2.0
|
|
|
|
import torch.nn.functional as F
|
|
from typing import Optional, Tuple
|
|
import torch
|
|
import time
|
|
|
|
from model.base import SelfForcingModel
|
|
import torch.distributed as dist
|
|
from utils.i2v_conditioning import (
|
|
_get_i2v_context_frames,
|
|
_i2v_loss_mask_like,
|
|
_overwrite_i2v_context,
|
|
_zero_i2v_context_timestep,
|
|
)
|
|
|
|
|
|
class DMD(SelfForcingModel):
|
|
def __init__(self, args, device):
|
|
"""
|
|
Initialize the DMD (Distribution Matching Distillation) module.
|
|
This class is self-contained and compute generator and fake score losses
|
|
in the forward pass.
|
|
"""
|
|
super().__init__(args, device)
|
|
self.num_frame_per_block = getattr(args, "num_frame_per_block", 1)
|
|
self.same_step_across_blocks = getattr(args, "same_step_across_blocks", True)
|
|
self.min_num_training_frames = getattr(args, "min_num_training_frames", 21)
|
|
self.num_training_frames = getattr(args, "num_training_frames", 21)
|
|
|
|
if self.num_frame_per_block > 1:
|
|
for diffusion_model in (self.generator, self.real_score, self.fake_score):
|
|
if hasattr(diffusion_model.model, "num_frame_per_block"):
|
|
diffusion_model.model.num_frame_per_block = self.num_frame_per_block
|
|
|
|
self.independent_first_frame = getattr(args, "independent_first_frame", False)
|
|
if self.independent_first_frame and not getattr(args, "i2v", False):
|
|
self.generator.model.independent_first_frame = True
|
|
if args.gradient_checkpointing:
|
|
self.generator.enable_gradient_checkpointing()
|
|
self.fake_score.enable_gradient_checkpointing()
|
|
|
|
# this will be init later with fsdp-wrapped modules
|
|
self.inference_pipeline: SelfForcingTrainingPipeline = None
|
|
|
|
# Step 2: Initialize all dmd hyperparameters
|
|
self.num_train_timestep = getattr(args, "num_train_timestep", 1000)
|
|
self.min_step = int(0.02 * self.num_train_timestep)
|
|
self.max_step = int(0.98 * self.num_train_timestep)
|
|
if hasattr(args, "real_guidance_scale"):
|
|
self.real_guidance_scale = args.real_guidance_scale
|
|
self.fake_guidance_scale = args.fake_guidance_scale
|
|
else:
|
|
self.real_guidance_scale = args.guidance_scale
|
|
self.fake_guidance_scale = 0.0
|
|
self.timestep_shift = getattr(args, "timestep_shift", 1.0)
|
|
self.ts_schedule = getattr(args, "ts_schedule", True)
|
|
self.ts_schedule_max = getattr(args, "ts_schedule_max", False)
|
|
self.min_score_timestep = getattr(args, "min_score_timestep", 0)
|
|
|
|
if getattr(self.scheduler, "alphas_cumprod", None) is not None:
|
|
self.scheduler.alphas_cumprod = self.scheduler.alphas_cumprod.to(device)
|
|
else:
|
|
self.scheduler.alphas_cumprod = None
|
|
|
|
@staticmethod
|
|
def _slice_block_cond_dict(cond_dict, batch_size, new_num_segments):
|
|
"""Slice a block-wise conditional dict to keep only the last `new_num_segments` segments."""
|
|
pe = cond_dict["prompt_embeds"]
|
|
orig_segs = pe.shape[0] // batch_size
|
|
if orig_segs > new_num_segments:
|
|
pe = pe.reshape(batch_size, orig_segs, *pe.shape[1:])[:, -new_num_segments:]
|
|
return {**cond_dict, "prompt_embeds": pe.reshape(batch_size * new_num_segments, *pe.shape[2:])}
|
|
return cond_dict
|
|
|
|
def _compute_kl_grad(
|
|
self, noisy_image_or_video: torch.Tensor,
|
|
estimated_clean_image_or_video: torch.Tensor,
|
|
timestep: torch.Tensor,
|
|
conditional_dict: dict, unconditional_dict: dict,
|
|
normalization: bool = True,
|
|
clean_x: Optional[torch.Tensor] = None
|
|
) -> Tuple[torch.Tensor, dict]:
|
|
"""
|
|
Compute the KL grad (eq 7 in https://arxiv.org/abs/2311.18828).
|
|
Input:
|
|
- noisy_image_or_video: a tensor with shape [B, F, C, H, W] where the number of frame is 1 for images.
|
|
- estimated_clean_image_or_video: a tensor with shape [B, F, C, H, W] representing the estimated clean image or video.
|
|
- timestep: a tensor with shape [B, F] containing the randomly generated timestep.
|
|
- conditional_dict: a dictionary containing the conditional information (e.g. text embeddings, image embeddings).
|
|
- unconditional_dict: a dictionary containing the unconditional information (e.g. null/negative text embeddings, null/negative image embeddings).
|
|
- normalization: a boolean indicating whether to normalize the gradient.
|
|
Output:
|
|
- kl_grad: a tensor representing the KL grad.
|
|
- kl_log_dict: a dictionary containing the intermediate tensors for logging.
|
|
"""
|
|
# Step 1: Compute the fake score
|
|
_, pred_fake_image_cond = self.fake_score(
|
|
noisy_image_or_video=noisy_image_or_video,
|
|
conditional_dict=conditional_dict,
|
|
timestep=timestep,
|
|
clean_x=clean_x
|
|
)
|
|
|
|
if self.fake_guidance_scale != 0.0:
|
|
_, pred_fake_image_uncond = self.fake_score(
|
|
noisy_image_or_video=noisy_image_or_video,
|
|
conditional_dict=unconditional_dict,
|
|
timestep=timestep,
|
|
clean_x=clean_x
|
|
)
|
|
pred_fake_image = pred_fake_image_cond + (
|
|
pred_fake_image_cond - pred_fake_image_uncond
|
|
) * self.fake_guidance_scale
|
|
else:
|
|
pred_fake_image = pred_fake_image_cond
|
|
|
|
# Step 2: Compute the real score
|
|
_, pred_real_image_cond = self.real_score(
|
|
noisy_image_or_video=noisy_image_or_video,
|
|
conditional_dict=conditional_dict,
|
|
timestep=timestep,
|
|
clean_x=clean_x
|
|
)
|
|
|
|
_, pred_real_image_uncond = self.real_score(
|
|
noisy_image_or_video=noisy_image_or_video,
|
|
conditional_dict=unconditional_dict,
|
|
timestep=timestep,
|
|
clean_x=clean_x
|
|
)
|
|
|
|
pred_real_image = pred_real_image_cond + (
|
|
pred_real_image_cond - pred_real_image_uncond
|
|
) * self.real_guidance_scale
|
|
|
|
# Step 3: Compute the DMD gradient (DMD paper eq. 7).
|
|
grad = (pred_fake_image - pred_real_image)
|
|
|
|
# NOTE: Changed the normalizer for causal teacher — per-block normalization
|
|
if normalization:
|
|
p_real = (estimated_clean_image_or_video - pred_real_image)
|
|
|
|
B, F, C, H, W = p_real.shape
|
|
if dist.get_rank() == 0:
|
|
print(f"p_real: {p_real.shape}")
|
|
if (
|
|
self.independent_first_frame
|
|
and not getattr(self.args, "i2v", False)
|
|
and (F - 1) % self.num_frame_per_block == 0
|
|
):
|
|
p_real_tail = p_real[:, 1:]
|
|
p_real_blocks = p_real_tail.view(
|
|
B,
|
|
(F - 1) // self.num_frame_per_block,
|
|
self.num_frame_per_block,
|
|
C,
|
|
H,
|
|
W,
|
|
)
|
|
normalizer_tail = torch.abs(p_real_blocks).mean(dim=[2, 3, 4, 5], keepdim=True)
|
|
normalizer = torch.ones_like(p_real)
|
|
normalizer[:, 1:] = normalizer_tail.expand_as(p_real_blocks).reshape(
|
|
B, F - 1, C, H, W
|
|
)
|
|
else:
|
|
p_real_blocks = p_real.view(B, F // self.num_frame_per_block, self.num_frame_per_block, C, H, W)
|
|
normalizer = torch.abs(p_real_blocks).mean(dim=[2, 3, 4, 5], keepdim=True)
|
|
normalizer = normalizer.expand_as(p_real_blocks).reshape(B, F, C, H, W)
|
|
|
|
|
|
grad = grad / normalizer
|
|
grad = torch.nan_to_num(grad)
|
|
|
|
return grad, {
|
|
"dmdtrain_gradient_norm": torch.mean(torch.abs(grad)).detach(),
|
|
"timestep": timestep.detach()
|
|
}
|
|
|
|
def compute_distribution_matching_loss(
|
|
self,
|
|
image_or_video: torch.Tensor,
|
|
conditional_dict: dict,
|
|
unconditional_dict: dict,
|
|
gradient_mask: Optional[torch.Tensor] = None,
|
|
denoised_timestep_from: int = 0,
|
|
denoised_timestep_to: int = 0,
|
|
clean_x: Optional[torch.Tensor] = None,
|
|
initial_latent: Optional[torch.Tensor] = None,
|
|
) -> Tuple[torch.Tensor, dict]:
|
|
"""
|
|
Compute the DMD loss (eq 7 in https://arxiv.org/abs/2311.18828).
|
|
Input:
|
|
- image_or_video: a tensor with shape [B, F, C, H, W] where the number of frame is 1 for images.
|
|
- conditional_dict: a dictionary containing the conditional information (e.g. text embeddings, image embeddings).
|
|
- unconditional_dict: a dictionary containing the unconditional information (e.g. null/negative text embeddings, null/negative image embeddings).
|
|
- gradient_mask: a boolean tensor with the same shape as image_or_video indicating which pixels to compute loss .
|
|
Output:
|
|
- dmd_loss: a scalar tensor representing the DMD loss.
|
|
- dmd_log_dict: a dictionary containing the intermediate tensors for logging.
|
|
"""
|
|
context_frames = _get_i2v_context_frames(image_or_video, initial_latent)
|
|
original_latent = _overwrite_i2v_context(
|
|
image_or_video, initial_latent, context_frames
|
|
)
|
|
if clean_x is not None:
|
|
clean_x = _overwrite_i2v_context(clean_x, initial_latent, context_frames)
|
|
|
|
batch_size, num_frame = image_or_video.shape[:2]
|
|
|
|
with torch.no_grad():
|
|
# Step 1: Randomly sample timestep based on the given schedule and corresponding noise
|
|
min_timestep = denoised_timestep_to if self.ts_schedule and denoised_timestep_to is not None else self.min_score_timestep
|
|
max_timestep = denoised_timestep_from if self.ts_schedule_max and denoised_timestep_from is not None else self.num_train_timestep
|
|
timestep = self._get_timestep(
|
|
min_timestep,
|
|
max_timestep,
|
|
batch_size,
|
|
num_frame,
|
|
self.num_frame_per_block,
|
|
# t2v keeps the original NVFP4 behavior (one shared timestep across
|
|
# all frames); i2v uses per-block timesteps.
|
|
uniform_timestep=not getattr(self.args, "i2v", False)
|
|
)
|
|
|
|
# TODO:should we change it to `timestep = self.scheduler.timesteps[timestep]`?
|
|
if self.timestep_shift > 1:
|
|
timestep = self.timestep_shift * \
|
|
(timestep / 1000) / \
|
|
(1 + (self.timestep_shift - 1) * (timestep / 1000)) * 1000
|
|
timestep = timestep.clamp(self.min_step, self.max_step)
|
|
timestep = _zero_i2v_context_timestep(timestep, context_frames)
|
|
|
|
noise = torch.randn_like(image_or_video)
|
|
if context_frames > 0:
|
|
noise[:, :context_frames] = 0
|
|
noisy_latent = self.scheduler.add_noise(
|
|
original_latent.flatten(0, 1),
|
|
noise.flatten(0, 1),
|
|
timestep.flatten(0, 1)
|
|
).detach().unflatten(0, (batch_size, num_frame))
|
|
noisy_latent = _overwrite_i2v_context(
|
|
noisy_latent, initial_latent, context_frames
|
|
)
|
|
|
|
# Step 2: Compute the KL grad
|
|
grad, dmd_log_dict = self._compute_kl_grad(
|
|
noisy_image_or_video=noisy_latent,
|
|
estimated_clean_image_or_video=original_latent,
|
|
timestep=timestep,
|
|
conditional_dict=conditional_dict,
|
|
unconditional_dict=unconditional_dict,
|
|
clean_x=clean_x
|
|
)
|
|
|
|
context_mask = _i2v_loss_mask_like(original_latent, context_frames)
|
|
if context_mask is not None:
|
|
gradient_mask = context_mask if gradient_mask is None else gradient_mask & context_mask
|
|
|
|
if gradient_mask is not None:
|
|
dmd_loss = 0.5 * F.mse_loss(original_latent.double(
|
|
)[gradient_mask], (original_latent.double() - grad.double()).detach()[gradient_mask], reduction="mean")
|
|
else:
|
|
dmd_loss = 0.5 * F.mse_loss(original_latent.double(
|
|
), (original_latent.double() - grad.double()).detach(), reduction="mean")
|
|
return dmd_loss, dmd_log_dict
|
|
|
|
def generator_loss(
|
|
self,
|
|
image_or_video_shape,
|
|
conditional_dict: dict,
|
|
unconditional_dict: dict,
|
|
clean_latent: torch.Tensor,
|
|
initial_latent: torch.Tensor = None
|
|
) -> Tuple[torch.Tensor, dict]:
|
|
"""
|
|
Generate image/videos from noise and compute the DMD loss.
|
|
The noisy input to the generator is backward simulated.
|
|
This removes the need of any datasets during distillation.
|
|
See Sec 4.5 of the DMD2 paper (https://arxiv.org/abs/2405.14867) for details.
|
|
Input:
|
|
- image_or_video_shape: a list containing the shape of the image or video [B, F, C, H, W].
|
|
- conditional_dict: a dictionary containing the conditional information (e.g. text embeddings, image embeddings).
|
|
- unconditional_dict: a dictionary containing the unconditional information (e.g. null/negative text embeddings, null/negative image embeddings).
|
|
- clean_latent: a tensor containing the clean latents [B, F, C, H, W]. Need to be passed when no backward simulation is used.
|
|
Output:
|
|
- loss: a scalar tensor representing the generator loss.
|
|
- generator_log_dict: a dictionary containing the intermediate tensors for logging.
|
|
"""
|
|
# Step 1: Unroll generator to obtain fake videos
|
|
slice_last_frames = getattr(self.args, "slice_last_frames", 21)
|
|
_t_gen_start = time.time()
|
|
num_gen_frames = image_or_video_shape[1]
|
|
sampled_noise = torch.randn(
|
|
[image_or_video_shape[0], num_gen_frames, *image_or_video_shape[2:]],
|
|
device=self.device, dtype=self.dtype)
|
|
pred_image, gradient_mask, denoised_timestep_from, denoised_timestep_to = self._run_generator(
|
|
image_or_video_shape=image_or_video_shape,
|
|
conditional_dict=conditional_dict,
|
|
initial_latent=initial_latent,
|
|
slice_last_frames=slice_last_frames,
|
|
noise=sampled_noise,
|
|
clean_latent=clean_latent
|
|
)
|
|
gen_time = time.time() - _t_gen_start
|
|
# Step 2: Compute the DMD loss
|
|
_t_loss_start = time.time()
|
|
if getattr(self.args, "teacher_forcing", False):
|
|
if getattr(self.args, "backward_simulation", True):
|
|
score_clean_x = pred_image.detach()
|
|
else:
|
|
score_clean_x = clean_latent
|
|
else:
|
|
score_clean_x = None
|
|
_bs = pred_image.shape[0]
|
|
_new_segs = pred_image.shape[1] // self.num_frame_per_block
|
|
if not getattr(self.args, "generator_is_causal", True):
|
|
_new_segs = 1
|
|
conditional_dict = self._slice_block_cond_dict(conditional_dict, _bs, _new_segs)
|
|
unconditional_dict = self._slice_block_cond_dict(unconditional_dict, _bs, _new_segs)
|
|
dmd_loss, dmd_log_dict = self.compute_distribution_matching_loss(
|
|
image_or_video=pred_image,
|
|
conditional_dict=conditional_dict,
|
|
unconditional_dict=unconditional_dict,
|
|
gradient_mask=gradient_mask,
|
|
denoised_timestep_from=denoised_timestep_from,
|
|
denoised_timestep_to=denoised_timestep_to,
|
|
clean_x=score_clean_x,
|
|
initial_latent=initial_latent if pred_image.shape[1] == image_or_video_shape[1] else None,
|
|
)
|
|
try:
|
|
loss_val = dmd_loss.item()
|
|
except Exception:
|
|
loss_val = float('nan')
|
|
loss_time = time.time() - _t_loss_start
|
|
|
|
dmd_log_dict.update({
|
|
"gen_time": gen_time,
|
|
"loss_time": loss_time
|
|
})
|
|
|
|
return dmd_loss, dmd_log_dict
|
|
|
|
def critic_loss(
|
|
self,
|
|
image_or_video_shape,
|
|
conditional_dict: dict,
|
|
unconditional_dict: dict,
|
|
clean_latent: torch.Tensor,
|
|
initial_latent: torch.Tensor = None
|
|
) -> Tuple[torch.Tensor, dict]:
|
|
"""
|
|
Generate image/videos from noise and train the critic with generated samples.
|
|
The noisy input to the generator is backward simulated.
|
|
This removes the need of any datasets during distillation.
|
|
See Sec 4.5 of the DMD2 paper (https://arxiv.org/abs/2405.14867) for details.
|
|
Input:
|
|
- image_or_video_shape: a list containing the shape of the image or video [B, F, C, H, W].
|
|
- conditional_dict: a dictionary containing the conditional information (e.g. text embeddings, image embeddings).
|
|
- unconditional_dict: a dictionary containing the unconditional information (e.g. null/negative text embeddings, null/negative image embeddings).
|
|
- clean_latent: a tensor containing the clean latents [B, F, C, H, W]. Need to be passed when no backward simulation is used.
|
|
Output:
|
|
- loss: a scalar tensor representing the generator loss.
|
|
- critic_log_dict: a dictionary containing the intermediate tensors for logging.
|
|
"""
|
|
slice_last_frames = getattr(self.args, "slice_last_frames", 21)
|
|
# Step 1: Run generator on backward simulated noisy input
|
|
_t_gen_start = time.time()
|
|
with torch.no_grad():
|
|
num_gen_frames = image_or_video_shape[1]
|
|
sampled_noise = torch.randn(
|
|
[image_or_video_shape[0], num_gen_frames, *image_or_video_shape[2:]],
|
|
device=self.device, dtype=self.dtype)
|
|
generated_image, _, denoised_timestep_from, denoised_timestep_to = self._run_generator(
|
|
image_or_video_shape=image_or_video_shape,
|
|
conditional_dict=conditional_dict,
|
|
initial_latent=initial_latent,
|
|
slice_last_frames=slice_last_frames,
|
|
noise=sampled_noise,
|
|
clean_latent=clean_latent
|
|
)
|
|
gen_time = time.time() - _t_gen_start
|
|
score_initial_latent = (
|
|
initial_latent
|
|
if initial_latent is not None and generated_image.shape[1] == image_or_video_shape[1]
|
|
else None
|
|
)
|
|
context_frames = _get_i2v_context_frames(generated_image, score_initial_latent)
|
|
batch_size, num_frame = generated_image.shape[:2]
|
|
|
|
_new_segs = num_frame // self.num_frame_per_block
|
|
if not getattr(self.args, "generator_is_causal", True):
|
|
_new_segs = 1
|
|
conditional_dict = self._slice_block_cond_dict(conditional_dict, batch_size, _new_segs)
|
|
|
|
if getattr(self.args, "teacher_forcing", False):
|
|
if getattr(self.args, "backward_simulation", True):
|
|
score_clean_x = generated_image
|
|
else:
|
|
score_clean_x = clean_latent
|
|
else:
|
|
score_clean_x = None
|
|
if score_clean_x is not None:
|
|
score_clean_x = _overwrite_i2v_context(
|
|
score_clean_x, score_initial_latent, context_frames
|
|
)
|
|
_t_loss_start = time.time()
|
|
|
|
# Step 2: Compute the fake prediction
|
|
min_timestep = denoised_timestep_to if self.ts_schedule and denoised_timestep_to is not None else self.min_score_timestep
|
|
max_timestep = denoised_timestep_from if self.ts_schedule_max and denoised_timestep_from is not None else self.num_train_timestep
|
|
critic_timestep = self._get_timestep(
|
|
min_timestep,
|
|
max_timestep,
|
|
batch_size,
|
|
num_frame,
|
|
self.num_frame_per_block,
|
|
# t2v keeps the original NVFP4 behavior (one shared timestep across
|
|
# all frames); i2v uses per-block timesteps.
|
|
uniform_timestep=not getattr(self.args, "i2v", False)
|
|
)
|
|
|
|
if self.timestep_shift > 1:
|
|
critic_timestep = self.timestep_shift * \
|
|
(critic_timestep / 1000) / (1 + (self.timestep_shift - 1) * (critic_timestep / 1000)) * 1000
|
|
|
|
critic_timestep = critic_timestep.clamp(self.min_step, self.max_step)
|
|
critic_timestep = _zero_i2v_context_timestep(critic_timestep, context_frames)
|
|
|
|
critic_noise = torch.randn_like(generated_image)
|
|
if context_frames > 0:
|
|
critic_noise[:, :context_frames] = 0
|
|
noisy_generated_image = self.scheduler.add_noise(
|
|
generated_image.flatten(0, 1),
|
|
critic_noise.flatten(0, 1),
|
|
critic_timestep.flatten(0, 1)
|
|
).unflatten(0, (batch_size, num_frame))
|
|
noisy_generated_image = _overwrite_i2v_context(
|
|
noisy_generated_image, score_initial_latent, context_frames
|
|
)
|
|
|
|
_, pred_fake_image = self.fake_score(
|
|
noisy_image_or_video=noisy_generated_image,
|
|
conditional_dict=conditional_dict,
|
|
timestep=critic_timestep,
|
|
clean_x=score_clean_x
|
|
)
|
|
|
|
# Step 3: Compute the denoising loss for the fake critic
|
|
if getattr(self.args, "denoising_loss_type", "flow") == "flow":
|
|
from utils.wan_5b_wrapper import WanDiffusionWrapper
|
|
flow_pred = WanDiffusionWrapper._convert_x0_to_flow_pred(
|
|
scheduler=self.scheduler,
|
|
x0_pred=pred_fake_image.flatten(0, 1),
|
|
xt=noisy_generated_image.flatten(0, 1),
|
|
timestep=critic_timestep.flatten(0, 1)
|
|
)
|
|
pred_fake_noise = None
|
|
else:
|
|
flow_pred = None
|
|
pred_fake_noise = self.scheduler.convert_x0_to_noise(
|
|
x0=pred_fake_image.flatten(0, 1),
|
|
xt=noisy_generated_image.flatten(0, 1),
|
|
timestep=critic_timestep.flatten(0, 1)
|
|
).unflatten(0, (batch_size, num_frame))
|
|
|
|
denoising_loss = self.denoising_loss_func(
|
|
x=generated_image.flatten(0, 1),
|
|
x_pred=pred_fake_image.flatten(0, 1),
|
|
noise=critic_noise.flatten(0, 1),
|
|
noise_pred=pred_fake_noise,
|
|
alphas_cumprod=self.scheduler.alphas_cumprod,
|
|
timestep=critic_timestep.flatten(0, 1),
|
|
gradient_mask=(
|
|
_i2v_loss_mask_like(generated_image, context_frames).flatten(0, 1)
|
|
if context_frames > 0 else None
|
|
),
|
|
flow_pred=flow_pred,
|
|
)
|
|
|
|
try:
|
|
loss_val = denoising_loss.item()
|
|
except Exception:
|
|
loss_val = float('nan')
|
|
loss_time = time.time() - _t_loss_start
|
|
# Step 5: Debugging Log
|
|
critic_log_dict = {
|
|
"critic_timestep": critic_timestep.detach(),
|
|
"gen_time": gen_time,
|
|
"loss_time": loss_time
|
|
}
|
|
|
|
return denoising_loss, critic_log_dict
|