File size: 7,320 Bytes
59630ba | 1 2 3 4 5 6 7 8 9 10 11 12 13 14 15 16 17 18 19 20 21 22 23 24 25 26 27 28 29 30 31 32 33 34 35 36 37 38 39 40 41 42 43 44 45 46 47 48 49 50 51 52 53 54 55 56 57 58 59 60 61 62 63 64 65 66 67 68 69 70 71 72 73 74 75 76 77 78 79 80 81 82 83 84 85 86 87 88 89 90 91 92 93 94 95 96 97 98 99 100 101 102 103 104 105 106 107 108 109 110 111 112 113 114 115 116 117 118 119 120 121 122 123 124 125 126 127 128 129 130 131 132 133 134 135 136 137 138 139 140 141 142 143 144 145 146 147 148 149 150 151 152 153 154 155 156 | from typing import Tuple
import torch
from model.base import BaseModel
from utils.wan_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:
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)
def _initialize_models(self, args, device=None):
self.generator = WanDiffusionWrapper(**getattr(args, "model_kwargs", {}), is_causal=True)
self.generator.model.requires_grad_(True)
# Text encoder is optional; skip when training without text conditioning (e.g. WorldMem)
if getattr(args, "use_text_encoder", True):
self.text_encoder = WanTextEncoder()
self.text_encoder.requires_grad_(False)
else:
self.text_encoder = None
self.vae = WanVAEWrapper()
self.vae.requires_grad_(False)
self.scheduler = self.generator.get_scheduler()
if device is not None:
self.scheduler.timesteps = self.scheduler.timesteps.to(device)
# Set up PRoPE for all attention layers when action encoder is enabled
model_kwargs = getattr(args, "model_kwargs", {})
if model_kwargs.get("use_action_encoder", False):
image_or_video_shape = getattr(args, "image_or_video_shape", [1, 21, 16, 60, 104])
num_frames = image_or_video_shape[1]
latent_h = image_or_video_shape[3]
latent_w = image_or_video_shape[4]
# patch_size for WAN is (1, 2, 2)
patches_y = latent_h // 2
patches_x = latent_w // 2
self.generator.model.setup_prope(
patches_x=patches_x,
patches_y=patches_y,
image_width=latent_w,
image_height=latent_h,
cameras=num_frames,
prope_max_time=getattr(args, "prope_max_time", 60),
rope_scale_pattern=getattr(args, "rope_scale_pattern", "rope"),
)
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.
"""
noise = torch.randn_like(clean_latent)
batch_size, num_frame = image_or_video_shape[:2]
# 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)
noisy_latents = self.scheduler.add_noise(
clean_latent.flatten(0, 1),
noise.flatten(0, 1),
timestep.flatten(0, 1)
).unflatten(0, (batch_size, num_frame))
training_target = self.scheduler.training_target(clean_latent, noise, timestep)
# 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
# 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())
loss_unweighted = torch.nn.functional.mse_loss(
flow_pred.float(), training_target.float(), reduction='none'
).mean(dim=(2, 3, 4)) # (B, F) unweighted per-frame loss
weight = self.scheduler.training_weight(timestep).unflatten(0, (batch_size, num_frame))
loss = (loss_unweighted * weight).mean() # scalar; used for backprop
log_dict = {
"x0": clean_latent.detach(),
"x0_pred": x0_pred.detach(),
"loss_per_frame": loss_unweighted.detach(), # (B, F) unweighted
"loss_unweighted": loss_unweighted.detach().mean(), # scalar unweighted
}
return loss, log_dict
|