chore: import upstream snapshot with attribution
This commit is contained in:
wehub-resource-sync
@@ -0,0 +1,6 @@
|
||||
from .dmd import DMD
|
||||
from .diffusion import CausalDiffusion
|
||||
__all__ = [
|
||||
"DMD",
|
||||
"CausalDiffusion",
|
||||
]
|
||||
+545
@@ -0,0 +1,545 @@
|
||||
# Adopted from https://github.com/guandeh17/Self-Forcing
|
||||
# SPDX-License-Identifier: Apache-2.0
|
||||
from typing import Tuple
|
||||
from einops import rearrange
|
||||
from torch import nn
|
||||
import torch.distributed as dist
|
||||
import torch
|
||||
import math
|
||||
|
||||
from pipeline import SelfForcingTrainingPipeline
|
||||
from utils.config import section_get
|
||||
from utils.loss import get_denoising_loss
|
||||
from utils.wan_5b_wrapper import WanDiffusionWrapper, WanTextEncoder, WanVAEWrapper
|
||||
|
||||
|
||||
def build_default_denoising_step_list(sampling_steps, num_train_timesteps=1000, shift=1.0, include_zero=True):
|
||||
sigmas = torch.linspace(1.0, 0.0, int(sampling_steps) + 1, dtype=torch.float32)[:-1]
|
||||
sigmas = shift * sigmas / (1 + (shift - 1) * sigmas)
|
||||
timesteps = (sigmas * num_train_timesteps).to(torch.long)
|
||||
if include_zero:
|
||||
timesteps = torch.cat([timesteps, torch.zeros(1, dtype=torch.long)])
|
||||
return timesteps
|
||||
|
||||
|
||||
class BaseModel(nn.Module):
|
||||
def __init__(self, args, device):
|
||||
super().__init__()
|
||||
print("args.model_kwargs.model_name", args.model_kwargs.model_name)
|
||||
self._initialize_models(args, device)
|
||||
|
||||
self.device = device
|
||||
self.args = args
|
||||
self.independent_first_frame = getattr(args, "independent_first_frame", False)
|
||||
self.dtype = torch.bfloat16 if args.mixed_precision else torch.float32
|
||||
self.denoising_step_list = None
|
||||
if getattr(args, "denoising_step_list", None) is not None:
|
||||
self.denoising_step_list = torch.tensor(args.denoising_step_list, dtype=torch.long)
|
||||
if getattr(args, "warp_denoising_step", False):
|
||||
timesteps = torch.cat((self.scheduler.timesteps.cpu(), torch.tensor([0], dtype=torch.float32)))
|
||||
self.denoising_step_list = timesteps[1000 - self.denoising_step_list]
|
||||
elif getattr(args, "sampling_steps", None) is not None:
|
||||
self.denoising_step_list = build_default_denoising_step_list(
|
||||
sampling_steps=args.sampling_steps,
|
||||
num_train_timesteps=getattr(args, "num_train_timestep", self.scheduler.num_train_timesteps),
|
||||
shift=getattr(args, "timestep_shift", self.scheduler.shift),
|
||||
include_zero=True,
|
||||
)
|
||||
|
||||
def _initialize_models(self, args, device):
|
||||
self.real_model_name = getattr(args, "real_name", "Wan2.2-TI2V-5B")
|
||||
self.fake_model_name = getattr(args, "fake_name", "Wan2.2-TI2V-5B")
|
||||
self.local_attn_size = section_get(
|
||||
args,
|
||||
"inference",
|
||||
"local_attn_size",
|
||||
getattr(args, "model_kwargs", {}).get("local_attn_size", -1),
|
||||
aliases=("inference_local_attn_size",),
|
||||
)
|
||||
all_causal = getattr(args, "all_causal", False)
|
||||
score_is_causal = all_causal
|
||||
|
||||
model_name = args.model_kwargs.get("model_name", "Wan2.2-TI2V-5B")
|
||||
if "5B" not in model_name:
|
||||
raise ValueError(f"Only Wan2.2-TI2V-5B is supported in this release, got {model_name}")
|
||||
if not dist.is_initialized() or dist.get_rank() == 0:
|
||||
tag = "all-causal 5B mode" if all_causal else "Wan2.2-TI2V-5B"
|
||||
print(f"Using {tag}")
|
||||
|
||||
# Generator
|
||||
generator_is_causal = getattr(args, "generator_is_causal", True)
|
||||
self.generator = WanDiffusionWrapper(**getattr(args, "model_kwargs", {}), is_causal=generator_is_causal)
|
||||
self.generator.model.requires_grad_(True)
|
||||
|
||||
# Real Score
|
||||
real_kwargs = getattr(args, "real_model_kwargs", {"model_name": self.real_model_name})
|
||||
self.real_score = WanDiffusionWrapper(**real_kwargs, is_causal=score_is_causal)
|
||||
self.real_score.model.requires_grad_(False)
|
||||
|
||||
# Fake Score
|
||||
fake_kwargs = getattr(args, "fake_model_kwargs", {"model_name": self.fake_model_name})
|
||||
self.fake_score = WanDiffusionWrapper(**fake_kwargs, is_causal=score_is_causal)
|
||||
self.fake_score.model.requires_grad_(True)
|
||||
|
||||
# Text Encoder & VAE
|
||||
self.text_encoder = WanTextEncoder()
|
||||
self.text_encoder.requires_grad_(False)
|
||||
|
||||
self.vae = WanVAEWrapper()
|
||||
self.vae.requires_grad_(False)
|
||||
|
||||
self.scheduler = self.generator.get_scheduler()
|
||||
self.scheduler.timesteps = self.scheduler.timesteps.to(device)
|
||||
|
||||
def _get_timestep(
|
||||
self,
|
||||
min_timestep: int,
|
||||
max_timestep: int,
|
||||
batch_size: int,
|
||||
num_frame: int,
|
||||
num_frame_per_block: int,
|
||||
uniform_timestep: bool = False
|
||||
) -> torch.Tensor:
|
||||
"""
|
||||
Randomly generate a timestep tensor based on the generator's task type. It uniformly samples a timestep
|
||||
from the range [min_timestep, max_timestep], and returns a tensor of shape [batch_size, num_frame].
|
||||
- If uniform_timestep, it will use the same timestep for all frames.
|
||||
- If not uniform_timestep, it will use a different timestep for each block.
|
||||
"""
|
||||
if uniform_timestep:
|
||||
timestep = torch.randint(
|
||||
min_timestep,
|
||||
max_timestep,
|
||||
[batch_size, 1],
|
||||
device=self.device,
|
||||
dtype=torch.long
|
||||
).repeat(1, num_frame)
|
||||
return timestep
|
||||
else:
|
||||
timestep = torch.randint(
|
||||
min_timestep,
|
||||
max_timestep,
|
||||
[batch_size, num_frame],
|
||||
device=self.device,
|
||||
dtype=torch.long
|
||||
)
|
||||
# make the noise level the same within every block
|
||||
if self.independent_first_frame and not getattr(self.args, "i2v", False):
|
||||
# the first frame is always kept the same
|
||||
timestep_from_second = timestep[:, 1:]
|
||||
timestep_from_second = timestep_from_second.reshape(
|
||||
timestep_from_second.shape[0], -1, num_frame_per_block)
|
||||
timestep_from_second[:, :, 1:] = timestep_from_second[:, :, 0:1]
|
||||
timestep_from_second = timestep_from_second.reshape(
|
||||
timestep_from_second.shape[0], -1)
|
||||
timestep = torch.cat([timestep[:, 0:1], timestep_from_second], dim=1)
|
||||
else:
|
||||
timestep = timestep.reshape(
|
||||
timestep.shape[0], -1, num_frame_per_block)
|
||||
timestep[:, :, 1:] = timestep[:, :, 0:1]
|
||||
timestep = timestep.reshape(timestep.shape[0], -1)
|
||||
return timestep
|
||||
|
||||
|
||||
class SelfForcingModel(BaseModel):
|
||||
def __init__(self, args, device):
|
||||
super().__init__(args, device)
|
||||
self.denoising_loss_func = get_denoising_loss(getattr(args, "denoising_loss_type", "flow"))()
|
||||
|
||||
def _run_generator(
|
||||
self,
|
||||
image_or_video_shape,
|
||||
conditional_dict: dict,
|
||||
initial_latent: torch.tensor = None,
|
||||
slice_last_frames: int = 21,
|
||||
noise=None,
|
||||
clean_latent: torch.Tensor = None,
|
||||
) -> Tuple[torch.Tensor, torch.Tensor]:
|
||||
"""
|
||||
Optionally simulate the generator's input from noise using backward simulation
|
||||
and then run the generator for one-step.
|
||||
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).
|
||||
- initial_latent: a tensor containing the initial latents [B, F, C, H, W].
|
||||
- slice_last_frames: number of frames to keep from the end.
|
||||
- noise: optional pre-sampled noise.
|
||||
- clean_latent: a tensor [B, F, C, H, W] for off-policy mode (backward_simulation=False).
|
||||
Output:
|
||||
- pred_image: a tensor with shape [B, F, C, H, W].
|
||||
- gradient_mask: boolean tensor or None.
|
||||
- denoised_timestep_from: int or None.
|
||||
- denoised_timestep_to: int or None.
|
||||
"""
|
||||
use_backward_simulation = getattr(self.args, "backward_simulation", True)
|
||||
|
||||
if use_backward_simulation:
|
||||
return self._run_generator_backward_simulation(
|
||||
image_or_video_shape=image_or_video_shape,
|
||||
conditional_dict=conditional_dict,
|
||||
initial_latent=initial_latent,
|
||||
slice_last_frames=slice_last_frames,
|
||||
noise=noise,
|
||||
)
|
||||
else:
|
||||
assert clean_latent is not None, "clean_latent is required when backward_simulation=False"
|
||||
return self._run_generator_off_policy(
|
||||
image_or_video_shape=image_or_video_shape,
|
||||
conditional_dict=conditional_dict,
|
||||
clean_latent=clean_latent,
|
||||
initial_latent=initial_latent,
|
||||
)
|
||||
|
||||
def _run_generator_off_policy(
|
||||
self,
|
||||
image_or_video_shape,
|
||||
conditional_dict: dict,
|
||||
clean_latent: torch.Tensor,
|
||||
initial_latent: torch.Tensor = None,
|
||||
) -> Tuple[torch.Tensor, torch.Tensor, int, int]:
|
||||
"""
|
||||
Off-policy generator: add noise to clean_latent at each timestep in
|
||||
denoising_step_list, randomly pick one, and run the generator for a
|
||||
single forward pass. Returns (pred_x0, gradient_mask,
|
||||
denoised_timestep_from, denoised_timestep_to).
|
||||
"""
|
||||
batch_size, num_frame = image_or_video_shape[:2]
|
||||
denoising_step_list = self.denoising_step_list.to(self.device)
|
||||
|
||||
# Build noisy versions at every timestep in the schedule
|
||||
simulated_noisy_input = []
|
||||
for ts in denoising_step_list:
|
||||
rand_noise = torch.randn_like(clean_latent)
|
||||
noisy_timestep = ts * torch.ones(
|
||||
[batch_size, num_frame], device=self.device, dtype=torch.long)
|
||||
|
||||
if ts.item() != 0:
|
||||
noisy_image = self.scheduler.add_noise(
|
||||
clean_latent.flatten(0, 1),
|
||||
rand_noise.flatten(0, 1),
|
||||
noisy_timestep.flatten(0, 1),
|
||||
).unflatten(0, (batch_size, num_frame))
|
||||
else:
|
||||
noisy_image = clean_latent
|
||||
simulated_noisy_input.append(noisy_image)
|
||||
|
||||
simulated_noisy_input = torch.stack(simulated_noisy_input, dim=1) # [B, T, F, C, H, W]
|
||||
|
||||
# Randomly sample a step index [B, F], uniform within each block
|
||||
num_steps = len(denoising_step_list)
|
||||
generator_is_causal = getattr(self.args, "generator_is_causal", True)
|
||||
if not generator_is_causal:
|
||||
# Bidirectional generator: all frames must share the same timestep
|
||||
index = torch.randint(0, num_steps, [batch_size, 1],
|
||||
device=self.device, dtype=torch.long).expand(-1, num_frame).contiguous()
|
||||
else:
|
||||
index = torch.randint(0, num_steps, [batch_size, num_frame],
|
||||
device=self.device, dtype=torch.long)
|
||||
# Make the index the same within every block
|
||||
if self.independent_first_frame and not getattr(self.args, "i2v", False):
|
||||
idx_rest = index[:, 1:]
|
||||
idx_rest = idx_rest.reshape(batch_size, -1, self.num_frame_per_block)
|
||||
idx_rest[:, :, 1:] = idx_rest[:, :, 0:1]
|
||||
index = torch.cat([index[:, :1], idx_rest.reshape(batch_size, -1)], dim=1)
|
||||
else:
|
||||
index = index.reshape(batch_size, -1, self.num_frame_per_block)
|
||||
index[:, :, 1:] = index[:, :, 0:1]
|
||||
index = index.reshape(batch_size, -1)
|
||||
|
||||
# Gather the noisy input corresponding to the sampled index
|
||||
noisy_input = torch.gather(
|
||||
simulated_noisy_input, dim=1,
|
||||
index=index.reshape(batch_size, 1, num_frame, 1, 1, 1).expand(
|
||||
-1, -1, -1, *image_or_video_shape[2:])
|
||||
).squeeze(1) # [B, F, C, H, W]
|
||||
|
||||
timestep = denoising_step_list[index] # [B, F]
|
||||
context_frames = int(initial_latent.shape[1]) if initial_latent is not None else 0
|
||||
if context_frames > 0:
|
||||
if context_frames >= num_frame:
|
||||
raise ValueError(
|
||||
f"initial_latent has {context_frames} frames but training clip has {num_frame}."
|
||||
)
|
||||
noisy_input[:, :context_frames] = initial_latent.to(
|
||||
device=noisy_input.device,
|
||||
dtype=noisy_input.dtype,
|
||||
)
|
||||
timestep[:, :context_frames] = 0
|
||||
|
||||
# Single forward pass through the generator
|
||||
_, pred_x0 = self.generator(
|
||||
noisy_image_or_video=noisy_input,
|
||||
conditional_dict=conditional_dict,
|
||||
timestep=timestep.float(),
|
||||
clean_x=clean_latent if getattr(self.args, "teacher_forcing", False) else None,
|
||||
)
|
||||
pred_x0 = pred_x0.to(self.dtype)
|
||||
|
||||
# Derive denoised_timestep_from / to from the sampled index for ts_schedule
|
||||
# Use the first batch element's first block index as the representative scalar
|
||||
rep_idx = index[0, 0].item()
|
||||
denoised_timestep_from = denoising_step_list[rep_idx].item()
|
||||
if rep_idx + 1 < num_steps:
|
||||
denoised_timestep_to = denoising_step_list[rep_idx + 1].item()
|
||||
else:
|
||||
denoised_timestep_to = 0
|
||||
|
||||
gradient_mask = None
|
||||
if context_frames > 0:
|
||||
pred_x0[:, :context_frames] = initial_latent.to(
|
||||
device=pred_x0.device,
|
||||
dtype=pred_x0.dtype,
|
||||
)
|
||||
gradient_mask = torch.ones_like(pred_x0, dtype=torch.bool)
|
||||
gradient_mask[:, :context_frames] = False
|
||||
return pred_x0, gradient_mask, denoised_timestep_from, denoised_timestep_to
|
||||
|
||||
def _run_generator_backward_simulation(
|
||||
self,
|
||||
image_or_video_shape,
|
||||
conditional_dict: dict,
|
||||
initial_latent: torch.tensor = None,
|
||||
slice_last_frames: int = 21,
|
||||
noise=None,
|
||||
) -> Tuple[torch.Tensor, torch.Tensor]:
|
||||
"""
|
||||
On-policy generator via backward simulation (original path).
|
||||
"""
|
||||
if initial_latent is not None:
|
||||
conditional_dict["initial_latent"] = initial_latent
|
||||
noise_shape = image_or_video_shape.copy()
|
||||
|
||||
separate_first_frame = self.independent_first_frame and not getattr(self.args, "i2v", False)
|
||||
min_num_frames = (self.min_num_training_frames - 1) if separate_first_frame else self.min_num_training_frames
|
||||
max_num_frames = self.num_training_frames - 1 if separate_first_frame else self.num_training_frames
|
||||
assert max_num_frames % self.num_frame_per_block == 0
|
||||
assert min_num_frames % self.num_frame_per_block == 0
|
||||
max_num_blocks = max_num_frames // self.num_frame_per_block
|
||||
min_num_blocks = min_num_frames // self.num_frame_per_block
|
||||
num_generated_blocks = torch.randint(min_num_blocks, max_num_blocks + 1, (1,), device=self.device)
|
||||
dist.broadcast(num_generated_blocks, src=0)
|
||||
num_generated_blocks = num_generated_blocks.item()
|
||||
num_generated_frames = num_generated_blocks * self.num_frame_per_block
|
||||
if separate_first_frame and initial_latent is None:
|
||||
num_generated_frames += 1
|
||||
min_num_frames += 1
|
||||
noise_shape[1] = num_generated_frames
|
||||
if noise is not None:
|
||||
noise = noise[:, :num_generated_frames]
|
||||
else:
|
||||
noise = torch.randn(noise_shape, device=self.device, dtype=self.dtype)
|
||||
|
||||
pred_image_or_video, denoised_timestep_from, denoised_timestep_to = self._consistency_backward_simulation(
|
||||
noise=noise,
|
||||
slice_last_frames=slice_last_frames,
|
||||
**conditional_dict,
|
||||
)
|
||||
|
||||
if slice_last_frames != -1 and pred_image_or_video.shape[1] > slice_last_frames:
|
||||
with torch.no_grad():
|
||||
if slice_last_frames > 1:
|
||||
latent_to_decode = pred_image_or_video[:, :-(slice_last_frames - 1), ...]
|
||||
else:
|
||||
latent_to_decode = pred_image_or_video
|
||||
pixels = self.vae.decode_to_pixel(latent_to_decode)
|
||||
frame = pixels[:, -1:, ...].to(self.dtype)
|
||||
frame = rearrange(frame, "b t c h w -> b c t h w")
|
||||
image_latent = self.vae.encode_to_latent(frame).to(self.dtype)
|
||||
if slice_last_frames > 1:
|
||||
last_frames = pred_image_or_video[:, -(slice_last_frames - 1):, ...]
|
||||
pred_image_or_video_sliced = torch.cat([image_latent, last_frames], dim=1)
|
||||
else:
|
||||
pred_image_or_video_sliced = image_latent
|
||||
if num_generated_frames != min_num_frames:
|
||||
gradient_mask = torch.ones_like(pred_image_or_video_sliced, dtype=torch.bool)
|
||||
if self.independent_first_frame:
|
||||
gradient_mask[:, :1] = False
|
||||
else:
|
||||
gradient_mask[:, :self.num_frame_per_block] = False
|
||||
else:
|
||||
gradient_mask = None
|
||||
else:
|
||||
pred_image_or_video_sliced = pred_image_or_video
|
||||
if num_generated_frames != min_num_frames:
|
||||
gradient_mask = torch.ones_like(pred_image_or_video_sliced, dtype=torch.bool)
|
||||
if self.independent_first_frame:
|
||||
gradient_mask[:, :1] = False
|
||||
else:
|
||||
gradient_mask[:, :self.num_frame_per_block] = False
|
||||
else:
|
||||
gradient_mask = None
|
||||
|
||||
pred_image_or_video_sliced = pred_image_or_video_sliced.to(self.dtype)
|
||||
return pred_image_or_video_sliced, gradient_mask, denoised_timestep_from, denoised_timestep_to
|
||||
|
||||
def _consistency_backward_simulation(
|
||||
self,
|
||||
noise: torch.Tensor,
|
||||
slice_last_frames: int = 21,
|
||||
**conditional_dict: dict
|
||||
) -> torch.Tensor:
|
||||
"""
|
||||
Simulate the generator's input from noise to avoid training/inference mismatch.
|
||||
See Sec 4.5 of the DMD2 paper (https://arxiv.org/abs/2405.14867) for details.
|
||||
Here we use the consistency sampler (https://arxiv.org/abs/2303.01469)
|
||||
Input:
|
||||
- noise: a tensor sampled from N(0, 1) 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).
|
||||
Output:
|
||||
- output: a tensor with shape [B, T, F, C, H, W].
|
||||
T is the total number of timesteps. output[0] is a pure noise and output[i] and i>0
|
||||
represents the x0 prediction at each timestep.
|
||||
"""
|
||||
generator_is_causal = getattr(self.args, "generator_is_causal", True)
|
||||
if not generator_is_causal:
|
||||
return self._bidirectional_backward_simulation(
|
||||
noise=noise,
|
||||
slice_last_frames=slice_last_frames,
|
||||
**conditional_dict,
|
||||
)
|
||||
|
||||
if self.inference_pipeline is None:
|
||||
self._initialize_inference_pipeline()
|
||||
|
||||
return self.inference_pipeline.inference_with_trajectory(
|
||||
noise=noise, **conditional_dict, slice_last_frames=slice_last_frames
|
||||
)
|
||||
|
||||
def _bidirectional_backward_simulation(
|
||||
self,
|
||||
noise: torch.Tensor,
|
||||
slice_last_frames: int = 21,
|
||||
**conditional_dict: dict
|
||||
) -> Tuple[torch.Tensor, int, int]:
|
||||
"""
|
||||
Backward simulation for bidirectional (non-causal) generator.
|
||||
All frames are processed at once at each denoising step — no KV cache,
|
||||
no block-by-block processing.
|
||||
"""
|
||||
from wan_5b.utils.fm_solvers_unipc import FlowUniPCMultistepScheduler
|
||||
|
||||
batch_size, num_frames = noise.shape[:2]
|
||||
|
||||
# Resolve single-segment prompt for bidirectional model
|
||||
prompt_embeds = conditional_dict["prompt_embeds"]
|
||||
num_segments = prompt_embeds.shape[0] // batch_size
|
||||
if num_segments > 1:
|
||||
prompt_embeds = prompt_embeds.reshape(
|
||||
batch_size, num_segments, *prompt_embeds.shape[1:])[:, 0]
|
||||
cond = {**conditional_dict, "prompt_embeds": prompt_embeds}
|
||||
|
||||
# Setup UniPC scheduler
|
||||
sampling_steps = getattr(self.args, "sampling_steps", None) or len(self.denoising_step_list)
|
||||
shift = self.scheduler.shift
|
||||
num_train_timesteps = self.scheduler.num_train_timesteps
|
||||
sample_scheduler = FlowUniPCMultistepScheduler(
|
||||
num_train_timesteps=num_train_timesteps, shift=1, use_dynamic_shifting=False)
|
||||
sample_scheduler.set_timesteps(sampling_steps, device=noise.device, shift=shift)
|
||||
unipc_timesteps = sample_scheduler.timesteps
|
||||
num_denoising_steps = len(unipc_timesteps)
|
||||
|
||||
# Pick a random exit step (synchronized across ranks)
|
||||
last_step_only = getattr(self.args, "last_step_only", False)
|
||||
if last_step_only:
|
||||
exit_step = num_denoising_steps - 1
|
||||
else:
|
||||
exit_step_t = torch.randint(0, num_denoising_steps, (1,), device=self.device)
|
||||
dist.broadcast(exit_step_t, src=0)
|
||||
exit_step = exit_step_t.item()
|
||||
|
||||
# Multi-step denoising loop (full-sequence bidirectional forward each step)
|
||||
latents = noise
|
||||
for index, t in enumerate(unipc_timesteps):
|
||||
timestep = t * torch.ones(
|
||||
[batch_size, num_frames], device=noise.device, dtype=torch.float32)
|
||||
|
||||
if index < exit_step:
|
||||
with torch.no_grad():
|
||||
flow_pred, _ = self.generator(
|
||||
noisy_image_or_video=latents,
|
||||
conditional_dict=cond,
|
||||
timestep=timestep,
|
||||
)
|
||||
latents = sample_scheduler.step(
|
||||
flow_pred, t, latents, return_dict=False)[0]
|
||||
else:
|
||||
# Exit step: forward with gradient
|
||||
flow_pred, denoised_pred = self.generator(
|
||||
noisy_image_or_video=latents,
|
||||
conditional_dict=cond,
|
||||
timestep=timestep,
|
||||
)
|
||||
break
|
||||
|
||||
# Compute denoised_timestep_from / to
|
||||
if exit_step == num_denoising_steps - 1:
|
||||
denoised_timestep_to = 0
|
||||
denoised_timestep_from = 1000 - torch.argmin(
|
||||
(self.scheduler.timesteps.cuda() - unipc_timesteps[exit_step].cuda()).abs(), dim=0).item()
|
||||
else:
|
||||
denoised_timestep_to = 1000 - torch.argmin(
|
||||
(self.scheduler.timesteps.cuda() - unipc_timesteps[exit_step + 1].cuda()).abs(), dim=0).item()
|
||||
denoised_timestep_from = 1000 - torch.argmin(
|
||||
(self.scheduler.timesteps.cuda() - unipc_timesteps[exit_step].cuda()).abs(), dim=0).item()
|
||||
|
||||
return denoised_pred, denoised_timestep_from, denoised_timestep_to
|
||||
|
||||
def _initialize_inference_pipeline(self):
|
||||
"""
|
||||
Lazy initialize the inference pipeline during the first backward simulation run.
|
||||
Here we encapsulate the inference code with a model-dependent outside function.
|
||||
We pass our FSDP-wrapped modules into the pipeline to save memory.
|
||||
"""
|
||||
local_attn_size = section_get(
|
||||
self.args,
|
||||
"inference",
|
||||
"local_attn_size",
|
||||
getattr(self.args, "model_kwargs", {}).get("local_attn_size", -1),
|
||||
aliases=("inference_local_attn_size",),
|
||||
)
|
||||
sink_size = section_get(
|
||||
self.args,
|
||||
"inference",
|
||||
"sink_size",
|
||||
getattr(self.args, "model_kwargs", {}).get("sink_size", 0),
|
||||
aliases=("inference_sink_size",),
|
||||
)
|
||||
multi_shot_sink = section_get(self.args, "inference", "multi_shot_sink", False)
|
||||
multi_shot_rope_offset = section_get(
|
||||
self.args,
|
||||
"inference",
|
||||
"multi_shot_rope_offset",
|
||||
0.0,
|
||||
)
|
||||
scene_cut_prefix = section_get(self.args, "inference", "scene_cut_prefix", "[SCENE_CUT]")
|
||||
slice_last_frames = getattr(self.args, "slice_last_frames", 21)
|
||||
# do not use self.num_training_frames, because it is changed by generator_loss and critic_loss
|
||||
num_training_frames = getattr(self.args, "num_training_frames")
|
||||
if local_attn_size == -1:
|
||||
kv_cache_size = num_training_frames
|
||||
else:
|
||||
kv_cache_size = min(local_attn_size + slice_last_frames, num_training_frames)
|
||||
frame_seq_length = math.prod(self.args.image_or_video_shape[-2:]) // 4
|
||||
self.inference_pipeline = SelfForcingTrainingPipeline(
|
||||
denoising_step_list=self.denoising_step_list,
|
||||
scheduler=self.scheduler,
|
||||
generator=self.generator,
|
||||
num_frame_per_block=self.num_frame_per_block,
|
||||
independent_first_frame=self.independent_first_frame,
|
||||
same_step_across_blocks=getattr(
|
||||
self.args, "same_step_across_blocks", getattr(self, "same_step_across_blocks", False)
|
||||
),
|
||||
last_step_only=getattr(self.args, "last_step_only", False),
|
||||
num_max_frames=kv_cache_size,
|
||||
context_noise=getattr(self.args, "context_noise", 0),
|
||||
sampling_steps=getattr(self.args, "sampling_steps", None),
|
||||
local_attn_size=local_attn_size,
|
||||
sink_size=sink_size,
|
||||
multi_shot_sink=multi_shot_sink,
|
||||
scene_cut_prefix=scene_cut_prefix,
|
||||
multi_shot_rope_offset=multi_shot_rope_offset,
|
||||
slice_last_frames=slice_last_frames,
|
||||
num_training_frames=num_training_frames,
|
||||
model_name=getattr(self.args, "model_kwargs", {}).get("model_name", None),
|
||||
frame_seq_length=frame_seq_length,
|
||||
)
|
||||
@@ -0,0 +1,574 @@
|
||||
# Adopted from https://github.com/guandeh17/Self-Forcing
|
||||
# SPDX-License-Identifier: Apache-2.0
|
||||
|
||||
from typing import Tuple
|
||||
import random
|
||||
import torch
|
||||
|
||||
from model.base import BaseModel
|
||||
from pipeline import CausalDiffusionInferencePipeline
|
||||
from utils.i2v_conditioning import _overwrite_i2v_context, _zero_i2v_context_timestep
|
||||
from utils.wan_5b_wrapper import WanDiffusionWrapper, WanTextEncoder, WanVAEWrapper
|
||||
|
||||
|
||||
class CausalDiffusion(BaseModel):
|
||||
def __init__(self, args, device):
|
||||
"""
|
||||
Initialize the Diffusion loss module.
|
||||
"""
|
||||
super().__init__(args, device)
|
||||
self.num_frame_per_block = getattr(args, "num_frame_per_block", 1)
|
||||
if self.num_frame_per_block > 1:
|
||||
self.generator.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()
|
||||
|
||||
# Step 2: Initialize all hyperparameters
|
||||
self.num_train_timestep = args.num_train_timestep
|
||||
self.min_step = int(0.02 * self.num_train_timestep)
|
||||
self.max_step = int(0.98 * self.num_train_timestep)
|
||||
self.guidance_scale = args.guidance_scale
|
||||
self.timestep_shift = getattr(args, "timestep_shift", 1.0)
|
||||
self.teacher_forcing = getattr(args, "teacher_forcing", False)
|
||||
# Noise augmentation in teacher forcing, we add small noise to clean context latents
|
||||
self.noise_augmentation_max_timestep = getattr(args, "noise_augmentation_max_timestep", 0)
|
||||
|
||||
self.args = args
|
||||
self.device = device
|
||||
self.inference_pipeline = None
|
||||
|
||||
# Error recycling (SVI-style error buffer)
|
||||
# When ``enable_position_bucketing`` is true, each rank holds a 2D
|
||||
# buffer ``(local_block_position × timestep)``. The pos dimension only
|
||||
# covers the LOCAL slice of the sequence this rank is responsible for
|
||||
# (no cross-SP-rank pos sharing — those positions are simply not
|
||||
# reachable by this rank during forward), so memory cost scales as
|
||||
# ``num_blocks_global / sp_size`` instead of ``num_blocks_global``.
|
||||
# ``global_block_offset`` is recorded for logging only.
|
||||
# During the first ``buffer_warmup_iter`` global steps, errors are
|
||||
# all-gathered across the DP group (ranks with the same SP rank but
|
||||
# different DP replicas), so each rank's local pos buckets fill up
|
||||
# ``dp_size`` × faster without any wasted bandwidth.
|
||||
self.error_buffer = None
|
||||
self.noise_error_buffer = None
|
||||
self.er_num_blocks = 0 # local; >0 means 2D position-bucketed
|
||||
self.er_block_offset = 0 # global block offset for THIS rank
|
||||
er_cfg = getattr(args, "error_recycling", None)
|
||||
if er_cfg is not None and getattr(er_cfg, "enabled", False):
|
||||
from utils.error_buffer import build_error_buffer
|
||||
cfg_dict = er_cfg if isinstance(er_cfg, dict) else dict(er_cfg)
|
||||
cfg_dict.setdefault("num_train_timesteps", self.num_train_timestep)
|
||||
sp_size = int(getattr(args, "sequence_parallel_size", 1) or 1)
|
||||
if cfg_dict.get("enable_position_bucketing", False):
|
||||
shape = list(getattr(args, "image_or_video_shape", [1, 0]))
|
||||
total_frames = int(shape[1]) if len(shape) > 1 else 0
|
||||
assert total_frames > 0 and self.num_frame_per_block > 0, (
|
||||
"enable_position_bucketing=true requires "
|
||||
"image_or_video_shape[1] and num_frame_per_block to be set."
|
||||
)
|
||||
num_blocks_global = total_frames // self.num_frame_per_block
|
||||
assert num_blocks_global % sp_size == 0, (
|
||||
f"num_blocks_global ({num_blocks_global}) must be divisible "
|
||||
f"by sequence_parallel_size ({sp_size})."
|
||||
)
|
||||
self.er_num_blocks = num_blocks_global // sp_size # local
|
||||
# Determine this rank's SP index → global block offset.
|
||||
if sp_size > 1:
|
||||
import torch.distributed as dist
|
||||
if dist.is_initialized():
|
||||
sp_rank = dist.get_rank() % sp_size
|
||||
else:
|
||||
sp_rank = 0
|
||||
else:
|
||||
sp_rank = 0
|
||||
self.er_block_offset = sp_rank * self.er_num_blocks
|
||||
|
||||
# Shard timestep buckets across SP ranks: each SP rank only
|
||||
# stores t_bucket % sp_size == sp_rank, cutting per-rank CPU
|
||||
# memory by ~sp_size. This uses the same SP dimension that
|
||||
# already splits positions in 2D mode, so both 1D and 2D
|
||||
# follow one save/load pattern (per sp_rank).
|
||||
import torch.distributed as dist
|
||||
if sp_size > 1 and dist.is_initialized():
|
||||
er_shard_rank = dist.get_rank() % sp_size
|
||||
er_shard_size = sp_size
|
||||
else:
|
||||
er_shard_rank = 0
|
||||
er_shard_size = 1
|
||||
|
||||
self.error_buffer = build_error_buffer(
|
||||
cfg_dict, num_blocks=self.er_num_blocks,
|
||||
global_block_offset=self.er_block_offset,
|
||||
shard_rank=er_shard_rank, shard_size=er_shard_size,
|
||||
)
|
||||
self.noise_error_buffer = build_error_buffer(
|
||||
cfg_dict, num_blocks=self.er_num_blocks,
|
||||
global_block_offset=self.er_block_offset,
|
||||
shard_rank=er_shard_rank, shard_size=er_shard_size,
|
||||
)
|
||||
self.er_context_inject_prob = float(cfg_dict.get("context_inject_prob", 0.9))
|
||||
self.er_latent_inject_prob = float(cfg_dict.get("latent_inject_prob", 0.0))
|
||||
self.er_noise_inject_prob = float(cfg_dict.get("noise_inject_prob", 0.0))
|
||||
self.er_clean_prob = float(cfg_dict.get("clean_prob", 0.0))
|
||||
self.er_clean_buffer_update_prob = float(cfg_dict.get("clean_buffer_update_prob", 0.1))
|
||||
self.er_start_step = int(cfg_dict.get("start_step", 0))
|
||||
self.er_buffer_warmup_iter = int(cfg_dict.get("buffer_warmup_iter", 50))
|
||||
self.er_skip_block_0 = bool(cfg_dict.get("skip_block_0", False))
|
||||
|
||||
def _initialize_models(self, args, device):
|
||||
model_name = getattr(args.model_kwargs, "model_name", "Wan2.2-TI2V-5B")
|
||||
if "5B" not in model_name:
|
||||
raise ValueError(f"Only Wan2.2-TI2V-5B is supported in this release, got {model_name}")
|
||||
self.generator = WanDiffusionWrapper(**getattr(args, "model_kwargs", {}), is_causal=True)
|
||||
self.generator.model.requires_grad_(True)
|
||||
|
||||
self.text_encoder = WanTextEncoder()
|
||||
self.text_encoder.requires_grad_(False)
|
||||
|
||||
self.vae = WanVAEWrapper()
|
||||
self.vae.requires_grad_(False)
|
||||
|
||||
self.scheduler = self.generator.get_scheduler()
|
||||
self.scheduler.timesteps = self.scheduler.timesteps.to(device)
|
||||
|
||||
def generator_loss(
|
||||
self,
|
||||
image_or_video_shape,
|
||||
conditional_dict: dict,
|
||||
unconditional_dict: dict,
|
||||
clean_latent: torch.Tensor,
|
||||
initial_latent: torch.Tensor = None,
|
||||
loss_mask: torch.Tensor = None,
|
||||
loss_mask_global_valid_count: torch.Tensor = None,
|
||||
global_step: int = 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.
|
||||
- loss_mask: optional tensor of shape [B, F] with 1.0 for valid frames and 0.0 for padded frames.
|
||||
Under Sequence Parallel this is already the local chunk.
|
||||
- loss_mask_global_valid_count: optional scalar tensor with the total valid count across all SP ranks.
|
||||
When provided (SP mode), used as the denominator instead of the local loss_mask.sum().
|
||||
- global_step: current training step, used for error recycling delayed start.
|
||||
Output:
|
||||
- loss: a scalar tensor representing the generator loss.
|
||||
- generator_log_dict: a dictionary containing the intermediate tensors for logging.
|
||||
"""
|
||||
batch_size, num_frame = image_or_video_shape[:2]
|
||||
|
||||
noise = torch.randn_like(clean_latent)
|
||||
# Step 2: Randomly sample a timestep and add noise to denoiser inputs
|
||||
index = self._get_timestep(
|
||||
0,
|
||||
self.scheduler.num_train_timesteps,
|
||||
image_or_video_shape[0],
|
||||
image_or_video_shape[1],
|
||||
self.num_frame_per_block,
|
||||
uniform_timestep=False
|
||||
)
|
||||
timestep = self.scheduler.timesteps[index].to(dtype=self.dtype, device=self.device)
|
||||
context_latent = (
|
||||
initial_latent
|
||||
if getattr(self.args, "i2v", False) and initial_latent is not None
|
||||
else None
|
||||
)
|
||||
context_frames = int(context_latent.shape[1]) if context_latent is not None else 0
|
||||
if context_frames > 0:
|
||||
if context_frames >= num_frame:
|
||||
raise ValueError(
|
||||
f"initial_latent has {context_frames} frames but training clip has {num_frame}."
|
||||
)
|
||||
timestep[:, :context_frames] = 0
|
||||
|
||||
# Step 2.5 & 3.5: Error recycling — clean_prob acts as a master switch.
|
||||
# When clean_prob fires, skip ALL error injection and use pristine input;
|
||||
# otherwise each injection type rolls its own probability and is gated
|
||||
# only by whether the corresponding buffer has any samples (SVI behavior).
|
||||
#
|
||||
# NOTE on rank-sync: random.random() below is INTENTIONALLY independent
|
||||
# across ranks. None of these decisions guard a collective call (we use
|
||||
# SVI's pattern of unconditional all_gather + local random replay in
|
||||
# Step 5), so per-rank divergence here only affects which slice of data
|
||||
# gets corrupted on which rank — perfectly safe under DP+SP, and matches
|
||||
# SVI's behavior exactly.
|
||||
er_latent_injected = False
|
||||
er_noise_injected = False
|
||||
er_injected = False
|
||||
er_use_clean = False
|
||||
er_ready = (
|
||||
self.error_buffer is not None
|
||||
and (global_step is None or global_step >= self.er_start_step)
|
||||
)
|
||||
if er_ready and self.er_clean_prob > 0 and random.random() < self.er_clean_prob:
|
||||
er_use_clean = True
|
||||
|
||||
# Noise error injection (SVI's noise_prob): corrupt noise input.
|
||||
# training_target is then computed with the corrupted noise so the
|
||||
# model learns to predict a self-correcting velocity (SVI Eq. logic).
|
||||
noise_for_train = noise
|
||||
if (
|
||||
er_ready and not er_use_clean
|
||||
and self.er_noise_inject_prob > 0
|
||||
and not self.noise_error_buffer.is_empty()
|
||||
and random.random() < self.er_noise_inject_prob
|
||||
):
|
||||
noise_for_train = self._inject_noise_error_buffer(
|
||||
noise, index, batch_size, num_frame
|
||||
)
|
||||
er_noise_injected = True
|
||||
|
||||
# Latent error injection (SVI's latent_prob): corrupt clean_latent
|
||||
# before noising. training_target keeps pointing to ORIGINAL clean_latent.
|
||||
clean_latent_for_noise = clean_latent
|
||||
if (
|
||||
er_ready and not er_use_clean
|
||||
and self.er_latent_inject_prob > 0
|
||||
and not self.error_buffer.is_empty()
|
||||
and random.random() < self.er_latent_inject_prob
|
||||
):
|
||||
clean_latent_for_noise = self._inject_latent_error_buffer(
|
||||
clean_latent, index, batch_size, num_frame
|
||||
)
|
||||
er_latent_injected = True
|
||||
|
||||
noisy_latents = self.scheduler.add_noise(
|
||||
clean_latent_for_noise.flatten(0, 1),
|
||||
noise_for_train.flatten(0, 1),
|
||||
timestep.flatten(0, 1)
|
||||
).unflatten(0, (batch_size, num_frame))
|
||||
training_target = self.scheduler.training_target(clean_latent, noise_for_train, timestep)
|
||||
if context_frames > 0:
|
||||
noisy_latents[:, :context_frames] = context_latent.to(
|
||||
device=noisy_latents.device,
|
||||
dtype=noisy_latents.dtype,
|
||||
)
|
||||
training_target[:, :context_frames] = 0
|
||||
|
||||
# Step 3: Noise augmentation, also add small noise to clean context latents
|
||||
if self.noise_augmentation_max_timestep > 0:
|
||||
index_clean_aug = self._get_timestep(
|
||||
0,
|
||||
self.noise_augmentation_max_timestep,
|
||||
image_or_video_shape[0],
|
||||
image_or_video_shape[1],
|
||||
self.num_frame_per_block,
|
||||
uniform_timestep=False
|
||||
)
|
||||
timestep_clean_aug = self.scheduler.timesteps[index_clean_aug].to(dtype=self.dtype, device=self.device)
|
||||
clean_latent_aug = self.scheduler.add_noise(
|
||||
clean_latent.flatten(0, 1),
|
||||
noise.flatten(0, 1),
|
||||
timestep_clean_aug.flatten(0, 1)
|
||||
).unflatten(0, (batch_size, num_frame))
|
||||
else:
|
||||
clean_latent_aug = clean_latent
|
||||
timestep_clean_aug = None
|
||||
|
||||
# Step 3.5: Error recycling — inject sampled errors into clean prefix.
|
||||
# 2D mode: per-position (random timestep). 1D mode: SVI global sampling.
|
||||
if (
|
||||
er_ready and not er_use_clean
|
||||
and not self.error_buffer.is_empty()
|
||||
and random.random() < self.er_context_inject_prob
|
||||
):
|
||||
clean_latent_aug = self._inject_error_buffer(
|
||||
clean_latent_aug, index, batch_size, num_frame
|
||||
)
|
||||
er_injected = True
|
||||
|
||||
if context_frames > 0:
|
||||
clean_latent_aug = _overwrite_i2v_context(
|
||||
clean_latent_aug, context_latent, context_frames
|
||||
)
|
||||
if timestep_clean_aug is not None:
|
||||
timestep_clean_aug = _zero_i2v_context_timestep(
|
||||
timestep_clean_aug, context_frames
|
||||
)
|
||||
|
||||
# Compute loss
|
||||
flow_pred, x0_pred = self.generator(
|
||||
noisy_image_or_video=noisy_latents,
|
||||
conditional_dict=conditional_dict,
|
||||
timestep=timestep,
|
||||
clean_x=clean_latent_aug if self.teacher_forcing else None,
|
||||
aug_t=timestep_clean_aug if self.teacher_forcing else None,
|
||||
)
|
||||
loss = torch.nn.functional.mse_loss(
|
||||
flow_pred.float(), training_target.float(), reduction='none'
|
||||
).mean(dim=(2, 3, 4))
|
||||
loss = loss * self.scheduler.training_weight(timestep).unflatten(0, (batch_size, num_frame))
|
||||
if context_frames > 0:
|
||||
if loss_mask is None:
|
||||
loss_mask = torch.ones(
|
||||
(batch_size, num_frame),
|
||||
device=loss.device,
|
||||
dtype=loss.dtype,
|
||||
)
|
||||
loss_mask[:, :context_frames] = 0
|
||||
if loss_mask is not None:
|
||||
loss = loss * loss_mask
|
||||
valid_count = loss_mask_global_valid_count if loss_mask_global_valid_count is not None else loss_mask.sum()
|
||||
loss = loss.sum() / valid_count.clamp(min=1.0)
|
||||
else:
|
||||
loss = loss.mean()
|
||||
|
||||
log_dict = {
|
||||
"x0": clean_latent.detach(),
|
||||
"x0_pred": x0_pred.detach()
|
||||
}
|
||||
|
||||
# Step 5: Store prediction errors into error buffer.
|
||||
#
|
||||
# SVI-style two-phase pattern (avoids the rank-divergent collective
|
||||
# deadlock that an ``if random.random() < p: all_gather()`` would
|
||||
# introduce):
|
||||
# PHASE A (collective, UNCONDITIONAL): every rank reaches the
|
||||
# all_gather call together so NCCL stays in sync.
|
||||
# PHASE B (local, GATED): each rank independently decides whether
|
||||
# to actually replay the gathered items into its buffer. The
|
||||
# gate uses random.random() per rank — divergence here is fine
|
||||
# because no further collective follows.
|
||||
if self.error_buffer is not None:
|
||||
with torch.no_grad():
|
||||
# latent error: x0_pred - clean_latent ≡ -σ(v_pred - v_gt) — used for context/latent injection
|
||||
latent_err = x0_pred.detach() - clean_latent.detach()
|
||||
# noise error: SVI definition is (1-σ)(v_pred - v_gt) so the buffer entry,
|
||||
# when later added directly to noise, equals ε_pred - ε_gt.
|
||||
sigma = self.scheduler.sigmas.to(flow_pred.device)[index].reshape(
|
||||
batch_size, num_frame, 1, 1, 1
|
||||
).to(flow_pred.dtype)
|
||||
noise_err = (flow_pred.detach() - training_target.detach()) * (1.0 - sigma)
|
||||
|
||||
use_distributed = (
|
||||
global_step is not None
|
||||
and global_step <= self.er_buffer_warmup_iter
|
||||
)
|
||||
|
||||
# === PHASE A: collective — runs on EVERY rank, no gating ===
|
||||
if use_distributed:
|
||||
lat_items = self._gather_errors_for_buffer(
|
||||
self.error_buffer, latent_err, index, batch_size, num_frame
|
||||
)
|
||||
noise_items = self._gather_errors_for_buffer(
|
||||
self.noise_error_buffer, noise_err, index, batch_size, num_frame
|
||||
)
|
||||
else:
|
||||
lat_items = self._collect_local_items(
|
||||
self.error_buffer, latent_err, index, batch_size, num_frame
|
||||
)
|
||||
noise_items = self._collect_local_items(
|
||||
self.noise_error_buffer, noise_err, index, batch_size, num_frame
|
||||
)
|
||||
|
||||
# === PHASE B: local replay — random.random() per-rank is OK ===
|
||||
# When the input was clean (low-error), only update buffer with
|
||||
# small probability to avoid flooding it with near-zero samples
|
||||
# (SVI: clean_buffer_update_prob).
|
||||
should_update = True
|
||||
if er_use_clean and random.random() >= self.er_clean_buffer_update_prob:
|
||||
should_update = False
|
||||
if should_update:
|
||||
self._apply_gathered_items(self.error_buffer, lat_items)
|
||||
self._apply_gathered_items(self.noise_error_buffer, noise_items)
|
||||
buf_stats = self.error_buffer.stats()
|
||||
noise_buf_stats = self.noise_error_buffer.stats()
|
||||
log_dict["er_total_added"] = buf_stats["total_added"]
|
||||
log_dict["er_filled_buckets"] = buf_stats["filled_buckets"]
|
||||
log_dict["er_total_entries"] = buf_stats["total_entries"]
|
||||
log_dict["er_noise_total_entries"] = noise_buf_stats["total_entries"]
|
||||
log_dict["er_injected"] = er_injected
|
||||
log_dict["er_latent_injected"] = er_latent_injected
|
||||
log_dict["er_noise_injected"] = er_noise_injected
|
||||
|
||||
return loss, log_dict
|
||||
|
||||
|
||||
def _inject_error_buffer(self, clean_latent_aug, index, batch_size, num_frame):
|
||||
"""Inject errors into the clean prefix (E_img).
|
||||
|
||||
2D (position-bucketed): the i-th LOCAL prefix block draws from
|
||||
``buckets[(i, *)]`` with a RANDOM timestep — the clean prefix is
|
||||
the product of full ODE integration so its accumulated error can
|
||||
come from any noise level, but its magnitude scales with the
|
||||
block's global position. Note ``skip_block_0`` is interpreted in
|
||||
the GLOBAL frame: only the very first SP rank may skip its block 0.
|
||||
|
||||
1D (timestep-bucketed): falls back to SVI ``sample_global``.
|
||||
"""
|
||||
block_size = self.num_frame_per_block
|
||||
num_blocks = num_frame // block_size
|
||||
result = clean_latent_aug.clone()
|
||||
for b in range(batch_size):
|
||||
for blk in range(num_blocks):
|
||||
if self.er_skip_block_0 and (self.er_block_offset + blk) == 0:
|
||||
continue
|
||||
if self.er_num_blocks > 0:
|
||||
err = self.error_buffer.sample_pos_any_t(
|
||||
blk, device=result.device, dtype=result.dtype
|
||||
)
|
||||
else:
|
||||
err = self.error_buffer.sample_global(
|
||||
device=result.device, dtype=result.dtype
|
||||
)
|
||||
if err is not None:
|
||||
start = blk * block_size
|
||||
end = start + block_size
|
||||
result[b, start:end] = result[b, start:end] + err
|
||||
return result
|
||||
|
||||
def _inject_latent_error_buffer(self, clean_latent, index, batch_size, num_frame):
|
||||
"""Inject errors into clean_latent before noising (E_vid).
|
||||
|
||||
Matches BOTH block_position (LOCAL) and timestep when the buffer is
|
||||
2D, else only timestep (SVI default).
|
||||
"""
|
||||
block_size = self.num_frame_per_block
|
||||
num_blocks = num_frame // block_size
|
||||
index_per_block = index[:, ::block_size]
|
||||
result = clean_latent.clone()
|
||||
for b in range(batch_size):
|
||||
for blk in range(num_blocks):
|
||||
t_idx = index_per_block[b, blk].item()
|
||||
pos = blk if self.er_num_blocks > 0 else None
|
||||
err = self.error_buffer.sample(
|
||||
t_idx, device=result.device, dtype=result.dtype,
|
||||
block_pos=pos,
|
||||
)
|
||||
if err is not None:
|
||||
start = blk * block_size
|
||||
end = start + block_size
|
||||
result[b, start:end] = result[b, start:end] + err
|
||||
return result
|
||||
|
||||
def _inject_noise_error_buffer(self, noise, index, batch_size, num_frame):
|
||||
"""Inject errors into the noise (E_noise).
|
||||
|
||||
Same matching strategy as ``_inject_latent_error_buffer`` but reads
|
||||
from the dedicated noise buffer.
|
||||
"""
|
||||
block_size = self.num_frame_per_block
|
||||
num_blocks = num_frame // block_size
|
||||
index_per_block = index[:, ::block_size]
|
||||
result = noise.clone()
|
||||
for b in range(batch_size):
|
||||
for blk in range(num_blocks):
|
||||
t_idx = index_per_block[b, blk].item()
|
||||
pos = blk if self.er_num_blocks > 0 else None
|
||||
err = self.noise_error_buffer.sample(
|
||||
t_idx, device=result.device, dtype=result.dtype,
|
||||
block_pos=pos,
|
||||
)
|
||||
if err is not None:
|
||||
start = blk * block_size
|
||||
end = start + block_size
|
||||
result[b, start:end] = result[b, start:end] + err
|
||||
return result
|
||||
|
||||
def _gather_errors_for_buffer(
|
||||
self, buffer, error, index, batch_size, num_frame
|
||||
):
|
||||
"""All-gather errors/timesteps across the appropriate group and return
|
||||
a list of ready-to-add ``(err_block, t_idx, pos_or_None)`` items.
|
||||
|
||||
★ This is a COLLECTIVE — every rank MUST reach this call together.
|
||||
The caller is responsible for invoking it unconditionally during the
|
||||
warmup window (just like SVI's ``all_gather`` outside the random
|
||||
``if`` blocks). Random decisions about whether to actually consume
|
||||
the returned items belong to ``_apply_gathered_items`` instead.
|
||||
|
||||
Group selection mirrors SVI's intent:
|
||||
* **2D (num_blocks > 0)** — DP group only. Other SP ranks' samples
|
||||
map to positions unreachable by this rank, so cross-SP gather
|
||||
wastes bandwidth.
|
||||
* **1D (num_blocks == 0)** — WORLD group (SVI default). Buckets
|
||||
are pos-agnostic so every rank's errors are valid samples.
|
||||
"""
|
||||
import torch.distributed as dist
|
||||
if not dist.is_initialized() or dist.get_world_size() <= 1:
|
||||
return self._collect_local_items(buffer, error, index, batch_size, num_frame)
|
||||
|
||||
if buffer.num_blocks > 0:
|
||||
from wan_5b.distributed.sp_training import get_data_parallel_group
|
||||
comm_group = get_data_parallel_group()
|
||||
if comm_group is None:
|
||||
return self._collect_local_items(buffer, error, index, batch_size, num_frame)
|
||||
comm_size = dist.get_world_size(comm_group)
|
||||
else:
|
||||
comm_group = None
|
||||
comm_size = dist.get_world_size()
|
||||
|
||||
if comm_size <= 1:
|
||||
return self._collect_local_items(buffer, error, index, batch_size, num_frame)
|
||||
|
||||
err_local = error.detach().contiguous()
|
||||
idx_local = index.detach().contiguous()
|
||||
err_list = [torch.empty_like(err_local) for _ in range(comm_size)]
|
||||
idx_list = [torch.empty_like(idx_local) for _ in range(comm_size)]
|
||||
if comm_group is None:
|
||||
dist.all_gather(err_list, err_local)
|
||||
dist.all_gather(idx_list, idx_local)
|
||||
else:
|
||||
dist.all_gather(err_list, err_local, group=comm_group)
|
||||
dist.all_gather(idx_list, idx_local, group=comm_group)
|
||||
|
||||
block_size = self.num_frame_per_block
|
||||
num_blocks = num_frame // block_size
|
||||
items = []
|
||||
for err_r, idx_r in zip(err_list, idx_list):
|
||||
idx_per_block = idx_r[:, ::block_size]
|
||||
err_blocks = err_r.reshape(
|
||||
batch_size, num_blocks, block_size, *err_r.shape[2:]
|
||||
)
|
||||
for b in range(batch_size):
|
||||
for blk in range(num_blocks):
|
||||
pos = blk if buffer.num_blocks > 0 else None
|
||||
items.append((err_blocks[b, blk], idx_per_block[b, blk].item(), pos))
|
||||
return items
|
||||
|
||||
def _collect_local_items(self, buffer, error, index, batch_size, num_frame):
|
||||
"""Same item-list format as ``_gather_errors_for_buffer`` but with no
|
||||
collective — used outside the warmup window or when distributed is off."""
|
||||
block_size = self.num_frame_per_block
|
||||
num_blocks = num_frame // block_size
|
||||
idx_per_block = index[:, ::block_size]
|
||||
error_blocks = error.reshape(
|
||||
batch_size, num_blocks, block_size, *error.shape[2:]
|
||||
)
|
||||
items = []
|
||||
for b in range(batch_size):
|
||||
for blk in range(num_blocks):
|
||||
pos = blk if buffer.num_blocks > 0 else None
|
||||
items.append((error_blocks[b, blk], idx_per_block[b, blk].item(), pos))
|
||||
return items
|
||||
|
||||
def _apply_gathered_items(self, buffer, items):
|
||||
"""Pure local: drop ``items`` into ``buffer``. No collective, no
|
||||
cross-rank coordination — each rank may invoke this independently
|
||||
(or skip it entirely) without risking a deadlock."""
|
||||
for err_block, t_idx, pos in items:
|
||||
buffer.add(err_block, t_idx, block_pos=pos)
|
||||
|
||||
def _initialize_inference_pipeline(self):
|
||||
"""
|
||||
Lazy initialize the inference pipeline during the first backward simulation run.
|
||||
Here we encapsulate the inference code with a model-dependent outside function.
|
||||
We pass our FSDP-wrapped modules into the pipeline to save memory.
|
||||
"""
|
||||
self.inference_pipeline = CausalDiffusionInferencePipeline(
|
||||
args=self.args,
|
||||
device=self.device,
|
||||
generator=self.generator,
|
||||
text_encoder=self.text_encoder,
|
||||
vae=self.vae
|
||||
)
|
||||
+494
@@ -0,0 +1,494 @@
|
||||
# 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
|
||||
Reference in New Issue
Block a user