500 lines
21 KiB
Python
500 lines
21 KiB
Python
# Modified from hhttps://github.com/OpenMOSS/MOVA/blob/main/mova/diffusion/pipelines/pipeline_mova.py
|
|
import math
|
|
import torch
|
|
import torch.nn as nn
|
|
from einops import rearrange
|
|
|
|
def sinusoidal_embedding_1d(dim, position):
|
|
sinusoid = torch.outer(position.type(torch.float64), torch.pow(
|
|
10000, -torch.arange(dim//2, dtype=torch.float64, device=position.device).div(dim//2)))
|
|
x = torch.cat([torch.cos(sinusoid), torch.sin(sinusoid)], dim=1)
|
|
return x.to(position.dtype)
|
|
|
|
|
|
class MOVAModel(nn.Module):
|
|
"""
|
|
MOVA helper class that encapsulates transformer, transformer_2, transformer_audio, and dual_tower_bridge.
|
|
Provides a clean forward interface similar to LTX2VideoTransformer3DModel.
|
|
This is NOT an nn.Module, just a helper class for organizing forward logic.
|
|
"""
|
|
def __init__(self, transformer, transformer_2, transformer_audio, dual_tower_bridge):
|
|
super().__init__()
|
|
self.transformer = transformer
|
|
self.transformer_2 = transformer_2
|
|
self.transformer_audio = transformer_audio
|
|
self.dual_tower_bridge = dual_tower_bridge
|
|
self.gradient_checkpointing = False
|
|
self.model_offload = False # Enable offloading unused models to CPU
|
|
|
|
@property
|
|
def dtype(self):
|
|
"""Return the dtype of the model (from first available transformer)."""
|
|
if self.transformer is not None:
|
|
return self.transformer.dtype
|
|
elif self.transformer_2 is not None:
|
|
return self.transformer_2.dtype
|
|
else:
|
|
raise AttributeError("MOVAModel has no available transformer to determine dtype")
|
|
|
|
@property
|
|
def config(self):
|
|
"""Return the config of the model (from first available transformer)."""
|
|
if self.transformer is not None:
|
|
return self.transformer.config
|
|
elif self.transformer_2 is not None:
|
|
return self.transformer_2.config
|
|
else:
|
|
raise AttributeError("MOVAModel has no available transformer to determine config")
|
|
|
|
def enable_gradient_checkpointing(self):
|
|
"""Enable gradient checkpointing for all sub-models to save memory."""
|
|
self.gradient_checkpointing = True
|
|
|
|
def disable_gradient_checkpointing(self):
|
|
"""Disable gradient checkpointing."""
|
|
self.gradient_checkpointing = False
|
|
|
|
def enable_model_offload(self):
|
|
"""Enable model offloading to save VRAM.
|
|
|
|
When enabled, only the active visual DiT (transformer or transformer_2)
|
|
and required components stay on GPU during forward pass.
|
|
"""
|
|
self.model_offload = True
|
|
|
|
def disable_model_offload(self):
|
|
"""Disable model offloading."""
|
|
self.model_offload = False
|
|
|
|
def set_module(self, module, module_name):
|
|
"""Disable model offloading."""
|
|
setattr(self, module_name, module)
|
|
|
|
def _move_to_device(self, model, device):
|
|
"""Helper to move model to device."""
|
|
if model is not None and model.device != torch.device(type="meta"):
|
|
model.to(device)
|
|
return model
|
|
|
|
def __call__(
|
|
self,
|
|
visual_latents: torch.Tensor,
|
|
audio_latents: torch.Tensor,
|
|
context: torch.Tensor,
|
|
timestep: torch.Tensor,
|
|
audio_timestep: torch.Tensor,
|
|
frame_rate: float,
|
|
use_low_noise_dit: bool = False,
|
|
):
|
|
"""
|
|
Forward pass for MOVA model.
|
|
|
|
Args:
|
|
visual_latents: [B, C_visual, T_v, H_v, W_v]
|
|
audio_latents: [B, C_audio, T_a]
|
|
context: [B, L_context, C_context]
|
|
timestep: [B] or scalar
|
|
audio_timestep: [B] or scalar
|
|
frame_rate: float
|
|
use_low_noise_dit: whether to use transformer (low noise, small t)
|
|
|
|
Returns:
|
|
visual_output: [B, C_visual, T_v, H_v, W_v]
|
|
audio_output: [B, C_audio, T_a]
|
|
"""
|
|
device = visual_latents.device
|
|
|
|
# Select which visual DiT to use
|
|
# Wan2.2 convention: transformer_2 = high-noise (large t), transformer = low-noise (small t)
|
|
active_visual_dit = self.transformer if use_low_noise_dit else self.transformer_2
|
|
inactive_visual_dit = self.transformer_2 if use_low_noise_dit else self.transformer
|
|
|
|
# Check if active model is available
|
|
if active_visual_dit is None:
|
|
raise ValueError(
|
|
f"Active visual DiT is None. use_low_noise_dit={use_low_noise_dit}. "
|
|
f"This may happen when training with boundary_type='low' or 'high'. "
|
|
f"Please check your training configuration."
|
|
)
|
|
|
|
# Model offloading: move inactive models to CPU to save VRAM
|
|
if self.model_offload:
|
|
# Move inactive visual DiT to CPU
|
|
if inactive_visual_dit is not None:
|
|
inactive_visual_dit.to('cpu')
|
|
torch.cuda.empty_cache()
|
|
|
|
# Move active visual DiT and transformer_audio to GPU
|
|
active_visual_dit = self._move_to_device(active_visual_dit, device)
|
|
self.transformer_audio = self._move_to_device(self.transformer_audio, device)
|
|
self.dual_tower_bridge = self._move_to_device(self.dual_tower_bridge, device)
|
|
else:
|
|
# No offload: just ensure models are on correct device
|
|
active_visual_dit = self._move_to_device(active_visual_dit, device)
|
|
self.transformer_audio = self._move_to_device(self.transformer_audio, device)
|
|
self.dual_tower_bridge = self._move_to_device(self.dual_tower_bridge, device)
|
|
|
|
output = self._forward_single_step(
|
|
visual_dit=active_visual_dit,
|
|
visual_latents=visual_latents,
|
|
audio_latents=audio_latents,
|
|
context=context,
|
|
timestep=timestep,
|
|
audio_timestep=audio_timestep,
|
|
frame_rate=frame_rate,
|
|
)
|
|
|
|
# Move active models back to CPU if offloading is enabled
|
|
if self.model_offload:
|
|
active_visual_dit.to('cpu')
|
|
self.transformer_audio.to('cpu')
|
|
self.dual_tower_bridge.to('cpu')
|
|
torch.cuda.empty_cache()
|
|
|
|
return output
|
|
|
|
def _forward_single_step(
|
|
self,
|
|
visual_dit,
|
|
visual_latents: torch.Tensor,
|
|
audio_latents: torch.Tensor,
|
|
context: torch.Tensor,
|
|
timestep: torch.Tensor,
|
|
audio_timestep: torch.Tensor,
|
|
frame_rate: float,
|
|
):
|
|
"""Single step forward pass."""
|
|
visual_x = visual_latents
|
|
audio_x = audio_latents
|
|
audio_context = visual_context = context
|
|
|
|
if audio_timestep is None:
|
|
audio_timestep = timestep
|
|
|
|
# Time embeddings
|
|
with torch.autocast("cuda", dtype=torch.float32):
|
|
visual_t = visual_dit.time_embedding(sinusoidal_embedding_1d(visual_dit.freq_dim, timestep))
|
|
visual_t_mod = visual_dit.time_projection(visual_t).unflatten(1, (6, visual_dit.dim))
|
|
|
|
audio_t = self.transformer_audio.time_embedding(sinusoidal_embedding_1d(self.transformer_audio.freq_dim, audio_timestep))
|
|
audio_t_mod = self.transformer_audio.time_projection(audio_t).unflatten(1, (6, self.transformer_audio.dim))
|
|
|
|
model_dtype = visual_dit.dtype
|
|
visual_t = visual_t.to(model_dtype)
|
|
visual_t_mod = visual_t_mod.to(model_dtype)
|
|
audio_t = audio_t.to(model_dtype)
|
|
audio_t_mod = audio_t_mod.to(model_dtype)
|
|
|
|
# Context embeddings
|
|
visual_context_emb = visual_dit.text_embedding(visual_context)
|
|
audio_context_emb = self.transformer_audio.text_embedding(audio_context)
|
|
|
|
visual_x = visual_latents.to(model_dtype)
|
|
audio_x = audio_latents.to(model_dtype)
|
|
|
|
# Visual patchify
|
|
visual_x = visual_x.contiguous(memory_format=torch.channels_last_3d)
|
|
visual_x = visual_dit.patch_embedding(visual_x)
|
|
grid_size = visual_x.shape[2:]
|
|
visual_x = rearrange(visual_x, 'b c f h w -> b (f h w) c').contiguous()
|
|
t, h, w = grid_size
|
|
|
|
# Audio patchify
|
|
audio_x = self.transformer_audio.patch_embedding(audio_x)
|
|
audio_grid_size = audio_x.shape[2:]
|
|
audio_x = rearrange(audio_x, 'b c f -> b f c').contiguous()
|
|
f = audio_grid_size[0]
|
|
|
|
# Audio freqs
|
|
audio_freqs = torch.cat(
|
|
[
|
|
self.transformer_audio.freqs[0][:f].view(f, -1).expand(f, -1),
|
|
self.transformer_audio.freqs[1][:f].view(f, -1).expand(f, -1),
|
|
self.transformer_audio.freqs[2][:f].view(f, -1).expand(f, -1),
|
|
],
|
|
dim=-1
|
|
).reshape(f, 1, -1).to(audio_x.device)
|
|
|
|
# Sequence parallel: chunk visual_x before blocks
|
|
# sp_world_size and sp_world_rank are also used in _forward_dual_tower_dit
|
|
self._sp_world_size = getattr(visual_dit, 'sp_world_size', 1)
|
|
self._sp_world_rank = getattr(visual_dit, 'sp_world_rank', 0)
|
|
if self._sp_world_size > 1:
|
|
# Pad sequence to be divisible by sp_world_size
|
|
seq_len = visual_x.shape[1]
|
|
padded_seq_len = int(math.ceil(seq_len / self._sp_world_size)) * self._sp_world_size
|
|
if padded_seq_len > seq_len:
|
|
visual_x = torch.cat([
|
|
visual_x,
|
|
visual_x.new_zeros(visual_x.shape[0], padded_seq_len - seq_len, visual_x.shape[2])
|
|
], dim=1)
|
|
# Chunk for sequence parallel
|
|
visual_x = torch.chunk(visual_x, self._sp_world_size, dim=1)[self._sp_world_rank]
|
|
|
|
# Prepare checkpointing utilities
|
|
if torch.is_grad_enabled() and self.gradient_checkpointing:
|
|
def create_custom_forward(module):
|
|
def custom_forward(*inputs):
|
|
return module(*inputs)
|
|
return custom_forward
|
|
|
|
# Forward through dual tower DiT blocks
|
|
visual_x, audio_x = self._forward_dual_tower_dit(
|
|
visual_dit=visual_dit,
|
|
visual_x=visual_x,
|
|
audio_x=audio_x,
|
|
visual_context=visual_context_emb,
|
|
audio_context=audio_context_emb,
|
|
visual_t_mod=visual_t_mod,
|
|
audio_t_mod=audio_t_mod,
|
|
grid_size=grid_size,
|
|
frame_rate=frame_rate,
|
|
)
|
|
|
|
# Sequence parallel: all_gather visual output after blocks
|
|
if self._sp_world_size > 1 and hasattr(visual_dit, 'all_gather') and visual_dit.all_gather is not None:
|
|
visual_x = visual_dit.all_gather(visual_x, dim=1)
|
|
|
|
# Visual head + unpatchify
|
|
if torch.is_grad_enabled() and self.gradient_checkpointing:
|
|
visual_output = torch.utils.checkpoint.checkpoint(
|
|
create_custom_forward(visual_dit.head), visual_x, visual_t, use_reentrant=False
|
|
)
|
|
else:
|
|
visual_output = visual_dit.head(visual_x, visual_t)
|
|
grid_sizes_tensor = torch.tensor([grid_size], dtype=torch.long, device=visual_output.device)
|
|
visual_output = visual_dit.unpatchify(visual_output, grid_sizes_tensor)
|
|
visual_output = visual_output[0].unsqueeze(0)
|
|
|
|
# Audio head + unpatchify
|
|
if torch.is_grad_enabled() and self.gradient_checkpointing:
|
|
audio_output = torch.utils.checkpoint.checkpoint(
|
|
create_custom_forward(self.transformer_audio.head), audio_x, audio_t, use_reentrant=False
|
|
)
|
|
else:
|
|
audio_output = self.transformer_audio.head(audio_x, audio_t)
|
|
audio_output = self.transformer_audio.unpatchify(audio_output, (f, ))
|
|
|
|
return visual_output, audio_output
|
|
|
|
def _forward_dual_tower_dit(
|
|
self,
|
|
visual_dit,
|
|
visual_x: torch.Tensor,
|
|
audio_x: torch.Tensor,
|
|
visual_context: torch.Tensor,
|
|
audio_context: torch.Tensor,
|
|
visual_t_mod: torch.Tensor,
|
|
audio_t_mod: torch.Tensor,
|
|
grid_size: tuple[int, int, int],
|
|
frame_rate: float,
|
|
condition_scale: float = 1.0,
|
|
a2v_condition_scale: float = None,
|
|
v2a_condition_scale: float = None,
|
|
):
|
|
"""Forward through dual tower DiT blocks with bridge."""
|
|
min_layers = min(len(visual_dit.blocks), len(self.transformer_audio.blocks))
|
|
visual_layers = len(visual_dit.blocks)
|
|
|
|
# Check if sequence parallel is enabled
|
|
sp_world_size = getattr(visual_dit, 'sp_world_size', 1)
|
|
sp_world_rank = getattr(visual_dit, 'sp_world_rank', 0)
|
|
sp_enabled = sp_world_size > 1 and hasattr(visual_dit, 'all_gather') and visual_dit.all_gather is not None
|
|
|
|
# Prepare checkpointing utilities
|
|
if torch.is_grad_enabled() and self.gradient_checkpointing:
|
|
def create_custom_forward(module):
|
|
def custom_forward(*inputs):
|
|
return module(*inputs)
|
|
return custom_forward
|
|
|
|
def create_custom_forward_bridge(module):
|
|
def custom_forward(layer_idx, visual_x, audio_x, x_freqs, y_freqs,
|
|
a2v_condition_scale, v2a_condition_scale,
|
|
condition_scale, video_grid_size):
|
|
return module(
|
|
layer_idx,
|
|
visual_x,
|
|
audio_x,
|
|
x_freqs=x_freqs,
|
|
y_freqs=y_freqs,
|
|
a2v_condition_scale=a2v_condition_scale,
|
|
v2a_condition_scale=v2a_condition_scale,
|
|
condition_scale=condition_scale,
|
|
video_grid_size=video_grid_size,
|
|
)
|
|
return custom_forward
|
|
|
|
ckpt_kwargs = {"use_reentrant": False}
|
|
|
|
# Prepare visual block parameters
|
|
t, h, w = grid_size
|
|
seq_len = t * h * w
|
|
visual_seq_lens = torch.tensor([seq_len], dtype=torch.long, device=visual_x.device)
|
|
visual_grid_sizes = torch.tensor([[t, h, w]], dtype=torch.long, device=visual_x.device)
|
|
visual_context_lens = None
|
|
visual_dtype = visual_x.dtype
|
|
wan_freqs = visual_dit.freqs.to(visual_x.device)
|
|
|
|
# Prepare audio block parameters
|
|
audio_f = audio_x.shape[1]
|
|
audio_seq_lens = torch.tensor([audio_f], dtype=torch.long, device=audio_x.device)
|
|
audio_grid_sizes = torch.tensor([[audio_f]], dtype=torch.long, device=audio_x.device)
|
|
audio_context_lens = None
|
|
audio_dtype = audio_x.dtype
|
|
audio_freqs_dit = torch.cat([
|
|
self.transformer_audio.freqs[0][:audio_f].view(audio_f, -1),
|
|
self.transformer_audio.freqs[1][:audio_f].view(audio_f, -1),
|
|
self.transformer_audio.freqs[2][:audio_f].view(audio_f, -1),
|
|
], dim=-1).reshape(audio_f, 1, -1).to(audio_x.device)
|
|
|
|
# Precompute cross-modal RoPE freqs
|
|
if self.dual_tower_bridge.apply_cross_rope:
|
|
(visual_rope_cos_sin, audio_rope_cos_sin) = self.dual_tower_bridge.build_aligned_freqs(
|
|
frame_rate=frame_rate,
|
|
grid_size=grid_size,
|
|
audio_steps=audio_x.shape[1],
|
|
device=visual_x.device,
|
|
dtype=visual_x.dtype,
|
|
)
|
|
else:
|
|
visual_rope_cos_sin = None
|
|
audio_rope_cos_sin = None
|
|
|
|
# Forward through blocks
|
|
for layer_idx in range(min_layers):
|
|
visual_block = visual_dit.blocks[layer_idx]
|
|
audio_block = self.transformer_audio.blocks[layer_idx]
|
|
|
|
# Cross-modal interaction via bridge with optional gradient checkpointing
|
|
# For sequence parallel: v2a (visual->audio) needs full visual sequence as key/value
|
|
# So we all_gather visual_x before bridge, then chunk it back after
|
|
needs_interaction = (
|
|
self.dual_tower_bridge.should_interact(layer_idx, 'a2v') or
|
|
self.dual_tower_bridge.should_interact(layer_idx, 'v2a')
|
|
)
|
|
|
|
if needs_interaction:
|
|
# Prepare visual_x for bridge: all_gather if sequence parallel is enabled
|
|
if sp_enabled:
|
|
visual_x_for_bridge = visual_dit.all_gather(visual_x, dim=1)
|
|
else:
|
|
visual_x_for_bridge = visual_x
|
|
|
|
if torch.is_grad_enabled() and self.gradient_checkpointing:
|
|
visual_x_out, audio_x = torch.utils.checkpoint.checkpoint(
|
|
create_custom_forward_bridge(self.dual_tower_bridge),
|
|
layer_idx,
|
|
visual_x_for_bridge,
|
|
audio_x,
|
|
visual_rope_cos_sin,
|
|
audio_rope_cos_sin,
|
|
a2v_condition_scale,
|
|
v2a_condition_scale,
|
|
condition_scale,
|
|
grid_size,
|
|
use_reentrant=False,
|
|
)
|
|
else:
|
|
visual_x_out, audio_x = self.dual_tower_bridge(
|
|
layer_idx,
|
|
visual_x_for_bridge,
|
|
audio_x,
|
|
x_freqs=visual_rope_cos_sin,
|
|
y_freqs=audio_rope_cos_sin,
|
|
a2v_condition_scale=a2v_condition_scale,
|
|
v2a_condition_scale=v2a_condition_scale,
|
|
condition_scale=condition_scale,
|
|
video_grid_size=grid_size,
|
|
)
|
|
|
|
# Chunk visual_x back to local rank if sequence parallel is enabled
|
|
# Bridge output visual_x might be modified (a2v direction), so always chunk
|
|
if sp_enabled:
|
|
visual_x = torch.chunk(visual_x_out, sp_world_size, dim=1)[sp_world_rank]
|
|
else:
|
|
visual_x = visual_x_out
|
|
|
|
# Visual block with optional gradient checkpointing
|
|
if torch.is_grad_enabled() and self.gradient_checkpointing:
|
|
visual_x = torch.utils.checkpoint.checkpoint(
|
|
create_custom_forward(visual_block),
|
|
visual_x,
|
|
visual_t_mod,
|
|
visual_seq_lens,
|
|
visual_grid_sizes,
|
|
wan_freqs,
|
|
visual_context,
|
|
visual_context_lens,
|
|
visual_dtype,
|
|
**ckpt_kwargs,
|
|
)
|
|
else:
|
|
visual_x = visual_block(
|
|
visual_x,
|
|
e=visual_t_mod,
|
|
seq_lens=visual_seq_lens,
|
|
grid_sizes=visual_grid_sizes,
|
|
freqs=wan_freqs,
|
|
context=visual_context,
|
|
context_lens=visual_context_lens,
|
|
dtype=visual_dtype,
|
|
)
|
|
|
|
# Audio block with optional gradient checkpointing
|
|
if torch.is_grad_enabled() and self.gradient_checkpointing:
|
|
audio_x = torch.utils.checkpoint.checkpoint(
|
|
create_custom_forward(audio_block),
|
|
audio_x,
|
|
audio_t_mod,
|
|
audio_seq_lens,
|
|
audio_grid_sizes,
|
|
audio_freqs_dit,
|
|
audio_context,
|
|
audio_context_lens,
|
|
audio_dtype,
|
|
**ckpt_kwargs,
|
|
)
|
|
else:
|
|
audio_x = audio_block(
|
|
audio_x,
|
|
e=audio_t_mod,
|
|
seq_lens=audio_seq_lens,
|
|
grid_sizes=audio_grid_sizes,
|
|
freqs=audio_freqs_dit,
|
|
context=audio_context,
|
|
context_lens=audio_context_lens,
|
|
dtype=audio_dtype,
|
|
)
|
|
|
|
# Forward remaining visual blocks
|
|
for layer_idx in range(min_layers, visual_layers):
|
|
visual_block = visual_dit.blocks[layer_idx]
|
|
|
|
if torch.is_grad_enabled() and self.gradient_checkpointing:
|
|
visual_x = torch.utils.checkpoint.checkpoint(
|
|
create_custom_forward(visual_block),
|
|
visual_x,
|
|
visual_t_mod,
|
|
visual_seq_lens,
|
|
visual_grid_sizes,
|
|
wan_freqs,
|
|
visual_context,
|
|
visual_context_lens,
|
|
visual_dtype,
|
|
**ckpt_kwargs,
|
|
)
|
|
else:
|
|
visual_x = visual_block(
|
|
visual_x,
|
|
e=visual_t_mod,
|
|
seq_lens=visual_seq_lens,
|
|
grid_sizes=visual_grid_sizes,
|
|
freqs=wan_freqs,
|
|
context=visual_context,
|
|
context_lens=visual_context_lens,
|
|
dtype=visual_dtype,
|
|
)
|
|
|
|
return visual_x, audio_x
|