[CI]Remove DreamID-Omni and MagiHuman support (#6357)
Signed-off-by: wangyu <410167048@qq.com> Signed-off-by: [Your Name] <your.email@example.com> Co-authored-by: Cursor <cursoragent@cursor.com>
This commit is contained in:
wangyuandCursor authored and GitHub committed
1 parent
40dcfdbb43
commit
aeb348d8f8
31 files changed
+26
-6957
No files matched your search
@@ -1,14 +0,0 @@
|
||||
# SPDX-License-Identifier: Apache-2.0
|
||||
# SPDX-FileCopyrightText: Copyright contributors to the vLLM project
|
||||
|
||||
__all__ = [
|
||||
"DreamIDOmniPipeline",
|
||||
]
|
||||
|
||||
|
||||
def __getattr__(name: str):
|
||||
if name == "DreamIDOmniPipeline":
|
||||
from .pipeline_dreamid_omni import DreamIDOmniPipeline
|
||||
|
||||
return DreamIDOmniPipeline
|
||||
raise AttributeError(f"module {__name__!r} has no attribute {name!r}")
|
||||
@@ -1,460 +0,0 @@
|
||||
import re
|
||||
from typing import TYPE_CHECKING
|
||||
|
||||
import torch
|
||||
import torch.nn as nn
|
||||
from cache_dit import ForwardPattern
|
||||
from vllm.distributed import get_tensor_model_parallel_world_size
|
||||
from vllm.logger import init_logger
|
||||
from vllm.model_executor.layers.linear import ColumnParallelLinear
|
||||
|
||||
from vllm_omni.diffusion.attention.layer import Attention
|
||||
from vllm_omni.diffusion.cache.cachedit import CacheDiTAdapterConfig
|
||||
|
||||
if TYPE_CHECKING:
|
||||
from vllm.model_executor.layers.quantization.base_config import QuantizationConfig
|
||||
|
||||
try:
|
||||
from dreamid_omni.modules.model import WanLayerNorm
|
||||
except ImportError:
|
||||
raise ImportError("Failed to import from dependency 'dreamid_omni'.")
|
||||
|
||||
from vllm_omni.diffusion.distributed.utils import get_local_device
|
||||
from vllm_omni.diffusion.models.dreamid_omni.wan2_2 import (
|
||||
DistributedRMSNorm,
|
||||
WanModel,
|
||||
rope_apply,
|
||||
)
|
||||
|
||||
logger = init_logger(__name__)
|
||||
|
||||
|
||||
class FusedBlock(nn.Module):
|
||||
"""Wrapper pairing a video block and audio block for layerwise offloading.
|
||||
|
||||
Registers both blocks as submodules so their parameters are visible to the offload hooks.
|
||||
"""
|
||||
|
||||
def __init__(
|
||||
self,
|
||||
vid_block: nn.Module,
|
||||
audio_block: nn.Module,
|
||||
device: torch.device,
|
||||
):
|
||||
super().__init__()
|
||||
self.vid_block = vid_block
|
||||
self.audio_block = audio_block
|
||||
self.device = device
|
||||
|
||||
def _cross_attention_forward(
|
||||
self,
|
||||
attn: Attention,
|
||||
cross_attn_block,
|
||||
src_seq,
|
||||
src_grid_sizes,
|
||||
src_freqs,
|
||||
target_seq,
|
||||
target_seq_lens,
|
||||
target_grid_sizes,
|
||||
target_freqs,
|
||||
context,
|
||||
context_lens,
|
||||
src_ref_lengths=None,
|
||||
src_freqs_scaling=None,
|
||||
target_ref_lengths=None,
|
||||
target_freqs_scaling=None,
|
||||
):
|
||||
b, n, d = src_seq.size(0), cross_attn_block.num_heads, cross_attn_block.head_dim
|
||||
if hasattr(cross_attn_block, "k_img"):
|
||||
q, k, v, k_img, v_img = cross_attn_block.qkv_fn(src_seq, context)
|
||||
else:
|
||||
q, k, v = cross_attn_block.qkv_fn(src_seq, context)
|
||||
k_img = v_img = None
|
||||
|
||||
x = attn(q, k, v)
|
||||
|
||||
if k_img is not None:
|
||||
img_x = attn(q, k_img, v_img)
|
||||
x = x + img_x
|
||||
|
||||
target_seq = cross_attn_block.pre_attn_norm_fusion(target_seq)
|
||||
k_target = cross_attn_block.norm_k_fusion(cross_attn_block.k_fusion(target_seq)).view(b, -1, n, d)
|
||||
v_target = cross_attn_block.v_fusion(target_seq).view(b, -1, n, d)
|
||||
|
||||
q = rope_apply(q, src_grid_sizes, src_freqs, ref_lengths=src_ref_lengths, freqs_scaling=src_freqs_scaling)
|
||||
k_target = rope_apply(
|
||||
k_target,
|
||||
target_grid_sizes,
|
||||
target_freqs,
|
||||
ref_lengths=target_ref_lengths,
|
||||
freqs_scaling=target_freqs_scaling,
|
||||
)
|
||||
|
||||
target_x = attn(q, k_target, v_target)
|
||||
|
||||
x = x + target_x
|
||||
x = x.flatten(2)
|
||||
x = cross_attn_block.o(x)
|
||||
return x
|
||||
|
||||
def _cross_attention_ffn_forward(
|
||||
self,
|
||||
attn: Attention,
|
||||
attn_block,
|
||||
src_seq,
|
||||
src_grid_sizes,
|
||||
src_freqs,
|
||||
target_seq,
|
||||
target_seq_lens,
|
||||
target_grid_sizes,
|
||||
target_freqs,
|
||||
context,
|
||||
context_lens,
|
||||
src_e,
|
||||
src_ref_lengths=None,
|
||||
src_freqs_scaling=None,
|
||||
target_ref_lengths=None,
|
||||
target_freqs_scaling=None,
|
||||
):
|
||||
src_seq = src_seq + self._cross_attention_forward(
|
||||
attn,
|
||||
attn_block.cross_attn,
|
||||
attn_block.norm3(src_seq),
|
||||
src_grid_sizes=src_grid_sizes,
|
||||
src_freqs=src_freqs,
|
||||
target_seq=target_seq,
|
||||
target_seq_lens=target_seq_lens,
|
||||
target_grid_sizes=target_grid_sizes,
|
||||
target_freqs=target_freqs,
|
||||
context=context,
|
||||
context_lens=context_lens,
|
||||
src_ref_lengths=src_ref_lengths,
|
||||
src_freqs_scaling=src_freqs_scaling,
|
||||
target_ref_lengths=target_ref_lengths,
|
||||
target_freqs_scaling=target_freqs_scaling,
|
||||
)
|
||||
y = attn_block.ffn(attn_block.norm2(src_seq).bfloat16() * (1 + src_e[4].squeeze(2)) + src_e[3].squeeze(2))
|
||||
with torch.autocast(device_type=self.device.type, dtype=torch.bfloat16, enabled=True):
|
||||
src_seq = src_seq + y * src_e[5].squeeze(2)
|
||||
return src_seq
|
||||
|
||||
def forward(
|
||||
self,
|
||||
hidden_states,
|
||||
encoder_hidden_states,
|
||||
attn: Attention,
|
||||
vid_e,
|
||||
vid_seq_lens,
|
||||
vid_grid_sizes,
|
||||
vid_freqs,
|
||||
vid_context,
|
||||
vid_context_lens,
|
||||
vid_ref_lengths,
|
||||
vid_freqs_scaling,
|
||||
audio_e,
|
||||
audio_seq_lens,
|
||||
audio_grid_sizes,
|
||||
audio_freqs,
|
||||
audio_context,
|
||||
audio_context_lens,
|
||||
audio_ref_lengths,
|
||||
audio_freqs_scaling,
|
||||
):
|
||||
vid = hidden_states
|
||||
audio = encoder_hidden_states
|
||||
vid_block = self.vid_block
|
||||
audio_block = self.audio_block
|
||||
|
||||
## audio modulation
|
||||
assert audio_e.dtype == torch.bfloat16
|
||||
assert len(audio_e.shape) == 4 and audio_e.size(2) == 6 and audio_e.shape[1] == audio.shape[1], (
|
||||
f"{audio_e.shape}, {audio.shape}"
|
||||
)
|
||||
with torch.autocast(device_type=self.device.type, dtype=torch.bfloat16, enabled=True):
|
||||
audio_e = audio_block.modulation(audio_e).chunk(6, dim=2)
|
||||
assert audio_e[0].dtype == torch.bfloat16
|
||||
|
||||
# audio self-attention
|
||||
audio_y = audio_block.self_attn(
|
||||
audio_block.norm1(audio).bfloat16() * (1 + audio_e[1].squeeze(2)) + audio_e[0].squeeze(2),
|
||||
audio_seq_lens,
|
||||
audio_grid_sizes,
|
||||
audio_freqs,
|
||||
ref_lengths=audio_ref_lengths,
|
||||
freqs_scaling=audio_freqs_scaling,
|
||||
)
|
||||
with torch.autocast(device_type=self.device.type, dtype=torch.bfloat16, enabled=True):
|
||||
audio = audio + audio_y * audio_e[2].squeeze(2)
|
||||
|
||||
## video modulation
|
||||
assert len(vid_e.shape) == 4 and vid_e.size(2) == 6 and vid_e.shape[1] == vid.shape[1], (
|
||||
f"{vid_e.shape}, {vid.shape}"
|
||||
)
|
||||
with torch.autocast(device_type=self.device.type, dtype=torch.bfloat16, enabled=True):
|
||||
vid_e = vid_block.modulation(vid_e).chunk(6, dim=2)
|
||||
|
||||
# video self-attention
|
||||
vid_y = vid_block.self_attn(
|
||||
vid_block.norm1(vid).bfloat16() * (1 + vid_e[1].squeeze(2)) + vid_e[0].squeeze(2),
|
||||
vid_seq_lens,
|
||||
vid_grid_sizes,
|
||||
vid_freqs,
|
||||
ref_lengths=vid_ref_lengths,
|
||||
)
|
||||
|
||||
with torch.autocast(device_type=self.device.type, dtype=torch.bfloat16, enabled=True):
|
||||
vid = vid + vid_y * vid_e[2].squeeze(2)
|
||||
|
||||
og_audio = audio
|
||||
|
||||
# audio cross-attention
|
||||
audio = self._cross_attention_ffn_forward(
|
||||
attn,
|
||||
audio_block,
|
||||
audio,
|
||||
audio_grid_sizes,
|
||||
audio_freqs,
|
||||
vid,
|
||||
vid_seq_lens,
|
||||
vid_grid_sizes,
|
||||
vid_freqs,
|
||||
audio_context,
|
||||
audio_context_lens,
|
||||
audio_e,
|
||||
src_ref_lengths=audio_ref_lengths,
|
||||
src_freqs_scaling=audio_freqs_scaling,
|
||||
target_ref_lengths=vid_ref_lengths,
|
||||
target_freqs_scaling=None,
|
||||
)
|
||||
|
||||
assert not torch.equal(og_audio, audio), "Audio should be changed after cross-attention!"
|
||||
|
||||
# video cross-attention
|
||||
vid = self._cross_attention_ffn_forward(
|
||||
attn,
|
||||
vid_block,
|
||||
vid,
|
||||
vid_grid_sizes,
|
||||
vid_freqs,
|
||||
og_audio,
|
||||
audio_seq_lens,
|
||||
audio_grid_sizes,
|
||||
audio_freqs,
|
||||
vid_context,
|
||||
vid_context_lens,
|
||||
vid_e,
|
||||
src_ref_lengths=vid_ref_lengths,
|
||||
src_freqs_scaling=None,
|
||||
target_ref_lengths=audio_ref_lengths,
|
||||
target_freqs_scaling=audio_freqs_scaling,
|
||||
)
|
||||
|
||||
return vid, audio
|
||||
|
||||
|
||||
class FusionModel(nn.Module):
|
||||
_cache_dit_adapter_config = CacheDiTAdapterConfig(
|
||||
block_forward_patterns={
|
||||
"fused_blocks": ForwardPattern.Pattern_0,
|
||||
},
|
||||
has_separate_cfg=True,
|
||||
)
|
||||
|
||||
_layerwise_offload_blocks_attrs = ["fused_blocks"]
|
||||
|
||||
packed_modules_mapping = {
|
||||
"to_qkv": ["q", "k", "v"],
|
||||
}
|
||||
|
||||
@staticmethod
|
||||
def _is_fused_block(name: str, module) -> bool:
|
||||
return "fused_blocks" in name and name.split(".")[-1].isdigit()
|
||||
|
||||
_hsdp_shard_conditions = [_is_fused_block]
|
||||
|
||||
def __init__(
|
||||
self,
|
||||
video_config=None,
|
||||
audio_config=None,
|
||||
*,
|
||||
quant_config: "QuantizationConfig | None" = None,
|
||||
):
|
||||
super().__init__()
|
||||
has_video = True
|
||||
has_audio = True
|
||||
self.device = get_local_device()
|
||||
if video_config is not None:
|
||||
self.video_model = WanModel(quant_config=quant_config, prefix="video_model", **video_config)
|
||||
else:
|
||||
has_video = False
|
||||
self.video_model = None
|
||||
logger.warning("No video model is provided!")
|
||||
|
||||
if audio_config is not None:
|
||||
self.audio_model = WanModel(quant_config=quant_config, prefix="audio_model", **audio_config)
|
||||
else:
|
||||
has_audio = False
|
||||
self.audio_model = None
|
||||
logger.warning("No audio model is provided!")
|
||||
|
||||
if has_video and has_audio:
|
||||
assert len(self.video_model.blocks) == len(self.audio_model.blocks)
|
||||
self.num_blocks = len(self.video_model.blocks)
|
||||
|
||||
self.inject_cross_attention_kv_projections(quant_config=quant_config)
|
||||
|
||||
tp_size = get_tensor_model_parallel_world_size()
|
||||
self.full_num_heads = self.video_model.num_heads
|
||||
self.num_heads = self.full_num_heads // tp_size
|
||||
self.head_dim = self.video_model.dim // self.video_model.num_heads
|
||||
self.attn = Attention(
|
||||
num_heads=self.num_heads,
|
||||
head_size=self.head_dim,
|
||||
num_kv_heads=self.num_heads,
|
||||
softmax_scale=1.0 / (self.head_dim**0.5),
|
||||
causal=False,
|
||||
)
|
||||
|
||||
if has_video and has_audio:
|
||||
self.fused_blocks = nn.ModuleList(
|
||||
[
|
||||
FusedBlock(
|
||||
self.video_model.blocks[i],
|
||||
self.audio_model.blocks[i],
|
||||
self.device,
|
||||
)
|
||||
for i in range(self.num_blocks)
|
||||
]
|
||||
)
|
||||
|
||||
def load_state_dict(self, state_dict, strict=True, assign=False):
|
||||
"""Remap checkpoints where blocks are stored under
|
||||
`video_model.blocks.N.*` / `audio_model.blocks.N.*` to the current
|
||||
`fused_blocks.N.vid_block.*` / `fused_blocks.N.audio_block.*`.
|
||||
"""
|
||||
needs_remap = any(re.match(r"^(video_model|audio_model)\.blocks\.\d+\.", k) for k in state_dict)
|
||||
if needs_remap:
|
||||
remapped = {}
|
||||
for k, v in state_dict.items():
|
||||
new_k = re.sub(r"^video_model\.blocks\.(\d+)\.", r"fused_blocks.\1.vid_block.", k)
|
||||
new_k = re.sub(r"^audio_model\.blocks\.(\d+)\.", r"fused_blocks.\1.audio_block.", new_k)
|
||||
remapped[new_k] = v
|
||||
state_dict = remapped
|
||||
|
||||
self._detach_blocks_from_backbones()
|
||||
|
||||
return super().load_state_dict(state_dict, strict=strict, assign=assign)
|
||||
|
||||
def inject_cross_attention_kv_projections(
|
||||
self,
|
||||
quant_config: "QuantizationConfig | None" = None,
|
||||
):
|
||||
def _make_kv_proj(dim: int, prefix: str) -> ColumnParallelLinear:
|
||||
# K/V projections for the cross-modal (a2v / v2a) attention path.
|
||||
# Inputs are full latents, so we quantize matched to self_attn.
|
||||
return ColumnParallelLinear(
|
||||
dim,
|
||||
dim,
|
||||
bias=True,
|
||||
gather_output=False,
|
||||
return_bias=False,
|
||||
quant_config=quant_config,
|
||||
prefix=prefix,
|
||||
)
|
||||
|
||||
tp_size = get_tensor_model_parallel_world_size()
|
||||
for i, vid_block in enumerate(self.video_model.blocks):
|
||||
cross = vid_block.cross_attn
|
||||
block_prefix = f"video_model.blocks.{i}.cross_attn"
|
||||
cross.k_fusion = _make_kv_proj(vid_block.dim, f"{block_prefix}.k_fusion")
|
||||
cross.v_fusion = _make_kv_proj(vid_block.dim, f"{block_prefix}.v_fusion")
|
||||
cross.pre_attn_norm_fusion = WanLayerNorm(vid_block.dim, elementwise_affine=True)
|
||||
cross.norm_k_fusion = (
|
||||
DistributedRMSNorm(vid_block.dim // tp_size, eps=1e-6) if vid_block.qk_norm else nn.Identity()
|
||||
)
|
||||
|
||||
for i, audio_block in enumerate(self.audio_model.blocks):
|
||||
cross = audio_block.cross_attn
|
||||
block_prefix = f"audio_model.blocks.{i}.cross_attn"
|
||||
cross.k_fusion = _make_kv_proj(audio_block.dim, f"{block_prefix}.k_fusion")
|
||||
cross.v_fusion = _make_kv_proj(audio_block.dim, f"{block_prefix}.v_fusion")
|
||||
cross.pre_attn_norm_fusion = WanLayerNorm(audio_block.dim, elementwise_affine=True)
|
||||
cross.norm_k_fusion = (
|
||||
DistributedRMSNorm(audio_block.dim // tp_size, eps=1e-6) if audio_block.qk_norm else nn.Identity()
|
||||
)
|
||||
|
||||
def _detach_blocks_from_backbones(self) -> None:
|
||||
"""Keep offloadable blocks owned only by a single place.
|
||||
|
||||
NOTE: This is a special workaround to support layerwise offloading.
|
||||
The model registers the same Wan blocks under both the video/audio
|
||||
backbones and `fused_blocks` which is a wrapper for unified blocks
|
||||
walking through. However, layerwise offloading will only consider
|
||||
`fused_blocks` as offloadable components and will materialize all
|
||||
other modules onto device, including the same blocks owned by both
|
||||
`fused_blocks` and `video_model` and `audio_model`.
|
||||
"""
|
||||
video_blocks = list(self.video_model.blocks)
|
||||
audio_blocks = list(self.audio_model.blocks)
|
||||
self.video_model._modules.pop("blocks", None)
|
||||
self.audio_model._modules.pop("blocks", None)
|
||||
self.video_model.blocks = tuple(video_blocks)
|
||||
self.audio_model.blocks = tuple(audio_blocks)
|
||||
|
||||
def merge_kwargs(self, vid_kwargs, audio_kwargs):
|
||||
"""
|
||||
keys in each kwarg:
|
||||
e
|
||||
seq_lens
|
||||
grid_sizes
|
||||
freqs
|
||||
context
|
||||
context_lens
|
||||
"""
|
||||
merged_kwargs = {}
|
||||
for key in vid_kwargs:
|
||||
merged_kwargs[f"vid_{key}"] = vid_kwargs[key]
|
||||
for key in audio_kwargs:
|
||||
merged_kwargs[f"audio_{key}"] = audio_kwargs[key]
|
||||
return merged_kwargs
|
||||
|
||||
def forward(
|
||||
self,
|
||||
vid,
|
||||
audio,
|
||||
t,
|
||||
vid_context,
|
||||
audio_context,
|
||||
vid_seq_len,
|
||||
audio_seq_len,
|
||||
ref_ip_lengths=None,
|
||||
ref_audio_lengths=None,
|
||||
slg_layer=False,
|
||||
freqs_scaling=None,
|
||||
):
|
||||
vid, vid_e, vid_kwargs = self.video_model.prepare_transformer_block_kwargs(
|
||||
x=vid, t=t, context=vid_context, seq_len=vid_seq_len, ref_lengths=ref_ip_lengths
|
||||
)
|
||||
|
||||
audio, audio_e, audio_kwargs = self.audio_model.prepare_transformer_block_kwargs(
|
||||
x=audio,
|
||||
t=t,
|
||||
context=audio_context,
|
||||
seq_len=audio_seq_len,
|
||||
ref_lengths=ref_audio_lengths,
|
||||
freqs_scaling=freqs_scaling,
|
||||
)
|
||||
|
||||
kwargs = self.merge_kwargs(vid_kwargs, audio_kwargs)
|
||||
|
||||
for fused_block in self.fused_blocks:
|
||||
vid, audio = fused_block(vid, audio, self.attn, **kwargs)
|
||||
|
||||
vid = self.video_model.post_transformer_block_out(vid, vid_kwargs["grid_sizes"], vid_e)
|
||||
audio = self.audio_model.post_transformer_block_out(audio, audio_kwargs["grid_sizes"], audio_e)
|
||||
|
||||
return vid, audio
|
||||
|
||||
def set_rope_params(self):
|
||||
self.video_model.set_rope_params()
|
||||
self.audio_model.set_rope_params()
|
||||
@@ -1,608 +0,0 @@
|
||||
# SPDX-License-Identifier: Apache-2.0
|
||||
# SPDX-FileCopyrightText: Copyright contributors to the vLLM project
|
||||
|
||||
import logging
|
||||
import math
|
||||
import os
|
||||
from collections.abc import Iterable
|
||||
from typing import ClassVar
|
||||
|
||||
import torch
|
||||
import torch.distributed
|
||||
from diffusers import FlowMatchEulerDiscreteScheduler
|
||||
from PIL import Image, ImageOps
|
||||
from torch import nn
|
||||
from torchvision.transforms import Compose, Normalize
|
||||
from tqdm import tqdm
|
||||
from vllm.distributed import (
|
||||
get_tensor_model_parallel_rank,
|
||||
get_tensor_model_parallel_world_size,
|
||||
)
|
||||
from vllm.model_executor.model_loader.weight_utils import default_weight_loader
|
||||
|
||||
from vllm_omni.diffusion.data import DiffusionOutput, OmniDiffusionConfig
|
||||
from vllm_omni.diffusion.distributed.cfg_parallel import CFGParallelMixin
|
||||
from vllm_omni.diffusion.distributed.utils import get_local_device
|
||||
from vllm_omni.diffusion.model_loader.diffusers_loader import DiffusersPipelineLoader
|
||||
from vllm_omni.diffusion.models.interface import SupportAudioInput, SupportImageInput, SupportsComponentDiscovery
|
||||
from vllm_omni.diffusion.worker.request_batch import DiffusionRequestBatch
|
||||
|
||||
try:
|
||||
from dreamid_omni.utils.divisible_crop import DivisibleCrop
|
||||
from dreamid_omni.utils.fm_solvers import FlowDPMSolverMultistepScheduler, get_sampling_sigmas, retrieve_timesteps
|
||||
from dreamid_omni.utils.fm_solvers_unipc import FlowUniPCMultistepScheduler
|
||||
from dreamid_omni.utils.model_loading_utils import (
|
||||
init_mmaudio_vae,
|
||||
init_text_model,
|
||||
init_wan_vae_2_2,
|
||||
)
|
||||
from dreamid_omni.utils.rearrange import Rearrange
|
||||
from dreamid_omni.utils.resize import NaResize
|
||||
except ImportError:
|
||||
raise ImportError("Failed to import from dependency 'dreamid_omni'.")
|
||||
from vllm_omni.diffusion.models.dreamid_omni.fusion import FusionModel
|
||||
|
||||
logger = logging.getLogger(__name__)
|
||||
|
||||
|
||||
def get_dreamid_omni_post_process_func(*args, **kwargs):
|
||||
def post_process(output):
|
||||
if isinstance(output, tuple) and len(output) == 2:
|
||||
video, audio = output
|
||||
return {
|
||||
"video": video,
|
||||
"audio": audio,
|
||||
"audio_sample_rate": 16000,
|
||||
"fps": 24,
|
||||
}
|
||||
return output
|
||||
|
||||
return post_process
|
||||
|
||||
|
||||
AUDIO_CONFIG = {
|
||||
"patch_size": [1],
|
||||
"model_type": "t2a",
|
||||
"dim": 3072,
|
||||
"ffn_dim": 14336,
|
||||
"freq_dim": 256,
|
||||
"num_heads": 24,
|
||||
"num_layers": 30,
|
||||
"in_dim": 20,
|
||||
"out_dim": 20,
|
||||
"text_len": 512,
|
||||
"window_size": [-1, -1],
|
||||
"qk_norm": True,
|
||||
"cross_attn_norm": True,
|
||||
"eps": 1e-6,
|
||||
"temporal_rope_scaling_factor": 0.19676,
|
||||
}
|
||||
|
||||
VIDEO_CONFIG = {
|
||||
"patch_size": [1, 2, 2],
|
||||
"model_type": "ti2v",
|
||||
"dim": 3072,
|
||||
"ffn_dim": 14336,
|
||||
"freq_dim": 256,
|
||||
"num_heads": 24,
|
||||
"num_layers": 30,
|
||||
"in_dim": 48,
|
||||
"out_dim": 48,
|
||||
"text_len": 512,
|
||||
"window_size": [-1, -1],
|
||||
"qk_norm": True,
|
||||
"cross_attn_norm": True,
|
||||
"eps": 1e-6,
|
||||
}
|
||||
|
||||
|
||||
class DreamIDOmniPipeline(
|
||||
nn.Module, CFGParallelMixin, SupportImageInput, SupportAudioInput, SupportsComponentDiscovery
|
||||
):
|
||||
"""DreamID-Omni pipeline for vLLM-Omni."""
|
||||
|
||||
_dit_modules: ClassVar[list[str]] = ["model"]
|
||||
_encoder_modules: ClassVar[list[str]] = ["text_model"]
|
||||
_vae_modules: ClassVar[list[str]] = ["vae_model_video", "vae_model_audio"]
|
||||
|
||||
def __init__(
|
||||
self,
|
||||
*,
|
||||
od_config: OmniDiffusionConfig,
|
||||
prefix: str = "",
|
||||
):
|
||||
super().__init__()
|
||||
self.od_config = od_config
|
||||
self.parallel_config = od_config.parallel_config
|
||||
|
||||
self.device = get_local_device()
|
||||
model = od_config.model
|
||||
self.target_dtype = od_config.dtype
|
||||
|
||||
# Init Models
|
||||
## Load VAEs
|
||||
vae_model_video = init_wan_vae_2_2(model, rank=self.device)
|
||||
vae_model_video.model.requires_grad_(False).eval()
|
||||
vae_model_video.model = vae_model_video.model.bfloat16()
|
||||
self.vae_model_video = vae_model_video
|
||||
self.vae = self.vae_model_video
|
||||
|
||||
vae_model_audio = init_mmaudio_vae(model, rank=self.device)
|
||||
vae_model_audio.requires_grad_(False).eval()
|
||||
self.vae_model_audio = vae_model_audio.bfloat16()
|
||||
|
||||
# Load T5 text model
|
||||
self.text_model = init_text_model(model, rank=self.device)
|
||||
self.text_encoder = self.text_model.model
|
||||
|
||||
# Fusion model — weights are loaded later via load_weights()
|
||||
quant_config = getattr(self.od_config, "quantization_config", None)
|
||||
self.model = FusionModel(VIDEO_CONFIG, AUDIO_CONFIG, quant_config=quant_config)
|
||||
self.transformer = self.model
|
||||
|
||||
fusion_path = self.od_config.tf_model_config.get("fusion", None)
|
||||
assert fusion_path is not None, "fusion checkpoint path is None in transformer config"
|
||||
fusion_subfolder = os.path.dirname(fusion_path) or None
|
||||
fusion_filename = os.path.basename(fusion_path)
|
||||
self.weights_sources = [
|
||||
DiffusersPipelineLoader.ComponentSource(
|
||||
model_or_path=model,
|
||||
subfolder=fusion_subfolder,
|
||||
revision=None,
|
||||
prefix="model.",
|
||||
allow_patterns_overrides=[fusion_filename],
|
||||
)
|
||||
]
|
||||
|
||||
# Fixed attributes, non-configurable
|
||||
self.audio_latent_channel = AUDIO_CONFIG.get("in_dim")
|
||||
self.video_latent_channel = VIDEO_CONFIG.get("in_dim")
|
||||
self.video_latent_length = 31
|
||||
self.audio_latent_length = 157
|
||||
self.target_area = 960 * 960
|
||||
|
||||
# Fixed attributes, CFG scales
|
||||
self.video_cfg_scale = 3.0
|
||||
self.video_ref_cfg_scale = 1.5
|
||||
self.audio_cfg_scale = 4.0
|
||||
self.audio_ref_cfg_scale = 2.0
|
||||
|
||||
# Schedulers will be set in forward
|
||||
self.scheduler_video = None
|
||||
self.scheduler_audio = None
|
||||
|
||||
def load_image_latent_ref_ip_video(
|
||||
self,
|
||||
images,
|
||||
audios,
|
||||
video_frame_height_width,
|
||||
):
|
||||
# Load size.
|
||||
patch_size = self.model.video_model.patch_size
|
||||
vae_stride = [4, 16, 16]
|
||||
|
||||
# Load image and video.
|
||||
ref_vae_latents = {
|
||||
"image": [],
|
||||
"audio": [],
|
||||
}
|
||||
video_h = video_frame_height_width[0]
|
||||
video_w = video_frame_height_width[1]
|
||||
ref_audio_lengths = []
|
||||
for img in images:
|
||||
# Calculate the required size to keep aspect ratio and fill the rest with padding.
|
||||
img_ratio = img.width / img.height
|
||||
target_ratio = video_w / video_h
|
||||
|
||||
if img_ratio > target_ratio: # Image is wider than target
|
||||
new_width = video_w
|
||||
new_height = int(new_width / img_ratio)
|
||||
else: # Image is taller than target
|
||||
new_height = video_h
|
||||
new_width = int(new_height * img_ratio)
|
||||
|
||||
img = img.resize((new_width, new_height), Image.Resampling.LANCZOS)
|
||||
|
||||
# Create a new image with the target size and place the resized image in the center
|
||||
delta_w = video_w - img.size[0]
|
||||
delta_h = video_h - img.size[1]
|
||||
padding = (delta_w // 2, delta_h // 2, delta_w - (delta_w // 2), delta_h - (delta_h // 2))
|
||||
new_img = ImageOps.expand(img, padding, fill=(255, 255, 255))
|
||||
|
||||
# Transform to tensor and normalize.
|
||||
image_transform = Compose(
|
||||
[
|
||||
NaResize(
|
||||
resolution=math.sqrt(
|
||||
video_frame_height_width[0] * video_frame_height_width[1]
|
||||
), # 256*448, 480*832
|
||||
mode="area",
|
||||
downsample_only=True,
|
||||
),
|
||||
DivisibleCrop((vae_stride[1] * patch_size[1], vae_stride[2] * patch_size[2])),
|
||||
Normalize(0.5, 0.5),
|
||||
Rearrange("t c h w -> c t h w"),
|
||||
]
|
||||
)
|
||||
new_img = image_transform([new_img])
|
||||
new_img = new_img.transpose(0, 1)
|
||||
new_img = new_img.to(self.device)
|
||||
new_img = new_img.to(self.target_dtype)
|
||||
|
||||
img_vae_latent = self.vae_model_video.wrapped_encode(new_img[:, :, None]).to(self.target_dtype).squeeze(0)
|
||||
ref_vae_latents["image"].append(img_vae_latent)
|
||||
|
||||
for audio_array in audios:
|
||||
audio_tensor = torch.from_numpy(audio_array).float().unsqueeze(0)
|
||||
# print(audio_tensor.shape)#torch.Size([1, 81280])
|
||||
audio_vae_latent = self.vae_model_audio.wrapped_encode(audio_tensor)
|
||||
# print(audio_vae_latent.shape)#torch.Size([1, 20, length])
|
||||
audio_length = audio_vae_latent.shape[2]
|
||||
ref_audio_lengths.append(audio_length)
|
||||
audio_vae_latent = audio_vae_latent.squeeze(0).transpose(0, 1)
|
||||
ref_vae_latents["audio"].append(audio_vae_latent)
|
||||
|
||||
ref_vae_latents["image"] = torch.cat(ref_vae_latents["image"], dim=1)
|
||||
|
||||
ref_vae_latents["audio"] = torch.cat(ref_vae_latents["audio"], dim=0)
|
||||
|
||||
return ref_vae_latents, ref_audio_lengths
|
||||
|
||||
def load_weights(self, weights: Iterable[tuple[str, torch.Tensor]]) -> set[str]:
|
||||
pipeline_prefix = "model."
|
||||
|
||||
qkv_stacked_mapping = [
|
||||
(".self_attn.to_qkv", ".self_attn.q", "q"),
|
||||
(".self_attn.to_qkv", ".self_attn.k", "k"),
|
||||
(".self_attn.to_qkv", ".self_attn.v", "v"),
|
||||
]
|
||||
|
||||
params_dict = dict(self.model.named_parameters())
|
||||
loaded_local: set[str] = set()
|
||||
|
||||
tp_size = get_tensor_model_parallel_world_size()
|
||||
tp_rank = get_tensor_model_parallel_rank() if tp_size > 1 else 0
|
||||
|
||||
def _maybe_shard_weight(weight: torch.Tensor, param: torch.Tensor) -> torch.Tensor:
|
||||
"""Slice a checkpoint tensor to this rank's TP shard when the
|
||||
param is a TP-local shape but the checkpoint stores the full one.
|
||||
"""
|
||||
if tp_size <= 1 or weight.shape == param.shape:
|
||||
return weight
|
||||
if weight.ndim == 1 and weight.numel() == param.numel() * tp_size:
|
||||
return weight.chunk(tp_size, dim=0)[tp_rank]
|
||||
if weight.ndim == 2:
|
||||
if weight.shape[0] == param.shape[0] * tp_size:
|
||||
return weight.chunk(tp_size, dim=0)[tp_rank]
|
||||
if weight.shape[1] == param.shape[1] * tp_size:
|
||||
return weight.chunk(tp_size, dim=1)[tp_rank]
|
||||
return weight
|
||||
|
||||
for full_name, loaded_weight in weights:
|
||||
if not full_name.startswith(pipeline_prefix):
|
||||
continue
|
||||
name = full_name[len(pipeline_prefix) :]
|
||||
|
||||
# Fused-QKV stacking: route checkpoint's separate
|
||||
# - ...self_attn.{q,k,v}: weights into the corresponding
|
||||
# - ...self_attn.to_qkv : via its sharded weight_loader.
|
||||
stacked_hit = False
|
||||
for fused_name, raw_name, shard_id in qkv_stacked_mapping:
|
||||
if raw_name not in name:
|
||||
continue
|
||||
stacked = name.replace(raw_name, fused_name)
|
||||
if stacked not in params_dict:
|
||||
continue
|
||||
param = params_dict[stacked]
|
||||
param.weight_loader(param, loaded_weight, shard_id)
|
||||
loaded_local.add(stacked)
|
||||
stacked_hit = True
|
||||
break
|
||||
if stacked_hit:
|
||||
continue
|
||||
|
||||
if name not in params_dict:
|
||||
logger.warning(f"Skipping DreamID-Omni weight {full_name} -- not found in model parameters")
|
||||
continue
|
||||
param = params_dict[name]
|
||||
weight_loader = getattr(param, "weight_loader", None)
|
||||
if weight_loader is not None:
|
||||
# Parallel Linears handle TP sharding internally
|
||||
weight_loader(param, loaded_weight)
|
||||
else:
|
||||
default_weight_loader(param, _maybe_shard_weight(loaded_weight, param))
|
||||
loaded_local.add(name)
|
||||
|
||||
# Detach must happen before `DiffusersPipelineLoader._process_weights_after_loading`
|
||||
# and before any forward pass that triggers layerwise-offload hooks.
|
||||
self.model._detach_blocks_from_backbones()
|
||||
|
||||
return {pipeline_prefix + k for k in loaded_local}
|
||||
|
||||
def get_scheduler_time_steps(self, sampling_steps, solver_name="unipc", device=0, shift=5.0):
|
||||
torch.manual_seed(4)
|
||||
|
||||
if solver_name == "unipc":
|
||||
sample_scheduler = FlowUniPCMultistepScheduler(
|
||||
num_train_timesteps=1000, shift=1, use_dynamic_shifting=False
|
||||
)
|
||||
sample_scheduler.set_timesteps(sampling_steps, device=device, shift=shift)
|
||||
timesteps = sample_scheduler.timesteps
|
||||
|
||||
elif solver_name == "dpm++":
|
||||
sample_scheduler = FlowDPMSolverMultistepScheduler(
|
||||
num_train_timesteps=1000, shift=1, use_dynamic_shifting=False
|
||||
)
|
||||
sampling_sigmas = get_sampling_sigmas(sampling_steps, shift=shift)
|
||||
timesteps, _ = retrieve_timesteps(sample_scheduler, device=device, sigmas=sampling_sigmas)
|
||||
|
||||
elif solver_name == "euler":
|
||||
sample_scheduler = FlowMatchEulerDiscreteScheduler(shift=shift)
|
||||
timesteps, sampling_steps = retrieve_timesteps(
|
||||
sample_scheduler,
|
||||
sampling_steps,
|
||||
device=device,
|
||||
)
|
||||
|
||||
else:
|
||||
raise NotImplementedError("Unsupported solver.")
|
||||
|
||||
return sample_scheduler, timesteps
|
||||
|
||||
def predict_noise(self, **kwargs):
|
||||
pred_vid, pred_audio = self.model(**kwargs)
|
||||
return (pred_vid[0], pred_audio[0])
|
||||
|
||||
def combine_multi_branch_cfg_noise(self, predictions, true_cfg_scale, cfg_normalize=False):
|
||||
vid_pos, audio_pos = predictions[0]
|
||||
vid_neg, audio_neg = predictions[1]
|
||||
vid_ip_neg, _ = predictions[2]
|
||||
_, refaudio_neg = predictions[3]
|
||||
|
||||
pred_video = (
|
||||
vid_neg
|
||||
+ true_cfg_scale["video_cfg_scale"] * (vid_pos - vid_neg)
|
||||
+ true_cfg_scale["video_ref_cfg_scale"] * (vid_pos - vid_ip_neg)
|
||||
)
|
||||
pred_audio = (
|
||||
audio_neg
|
||||
+ true_cfg_scale["audio_cfg_scale"] * (audio_pos - audio_neg)
|
||||
+ true_cfg_scale["audio_ref_cfg_scale"] * (audio_pos - refaudio_neg)
|
||||
)
|
||||
return (pred_video, pred_audio)
|
||||
|
||||
def diffuse(
|
||||
self,
|
||||
video_noise: torch.Tensor,
|
||||
audio_noise: torch.Tensor,
|
||||
latents_ref_image: torch.Tensor,
|
||||
latents_ref_audio: torch.Tensor,
|
||||
timesteps_video: torch.Tensor,
|
||||
timesteps_audio: torch.Tensor,
|
||||
text_embeddings_video_pos: torch.Tensor,
|
||||
text_embeddings_video_neg: torch.Tensor,
|
||||
text_embeddings_audio_pos: torch.Tensor,
|
||||
text_embeddings_audio_neg: torch.Tensor,
|
||||
max_seq_len_video: int,
|
||||
max_seq_len_audio: int,
|
||||
freqs_scaling_tensor: torch.Tensor,
|
||||
ref_ip_num: int,
|
||||
ref_audio_length: int,
|
||||
ref_audio_lengths: list,
|
||||
scheduler_video,
|
||||
scheduler_audio,
|
||||
) -> tuple[torch.Tensor, torch.Tensor]:
|
||||
"""
|
||||
Diffusion loop with CFG parallel support for DreamID-Omni.
|
||||
"""
|
||||
for i, (t_v, t_a) in tqdm(enumerate(zip(timesteps_video, timesteps_audio))):
|
||||
timestep_input = torch.full((1,), t_v, device=self.device)
|
||||
|
||||
model_input_video = torch.cat([video_noise[:, :-ref_ip_num], latents_ref_image], dim=1)
|
||||
model_input_video_neg = torch.cat(
|
||||
[video_noise[:, :-ref_ip_num], torch.zeros_like(latents_ref_image)], dim=1
|
||||
)
|
||||
|
||||
model_input_audio = torch.cat([audio_noise[:-ref_audio_length, :], latents_ref_audio], dim=0)
|
||||
|
||||
model_input_audio_neg = torch.cat(
|
||||
[audio_noise[:-ref_audio_length, :], torch.zeros_like(latents_ref_audio)], dim=0
|
||||
)
|
||||
ref_ip_lengths = [ref_ip_num]
|
||||
|
||||
common_args = {
|
||||
"vid_seq_len": max_seq_len_video,
|
||||
"audio_seq_len": max_seq_len_audio,
|
||||
"freqs_scaling": freqs_scaling_tensor,
|
||||
"ref_ip_lengths": [ref_ip_lengths],
|
||||
"ref_audio_lengths": [ref_audio_lengths],
|
||||
}
|
||||
pos_args = {
|
||||
**common_args,
|
||||
"audio_context": [text_embeddings_audio_pos],
|
||||
"vid_context": [text_embeddings_video_pos],
|
||||
}
|
||||
neg_args = {
|
||||
**common_args,
|
||||
"audio_context": [text_embeddings_audio_neg],
|
||||
"vid_context": [text_embeddings_video_neg],
|
||||
}
|
||||
|
||||
branches_kwargs = [
|
||||
{"vid": [model_input_video], "audio": [model_input_audio], "t": timestep_input, **pos_args},
|
||||
{"vid": [model_input_video], "audio": [model_input_audio], "t": timestep_input, **neg_args},
|
||||
{"vid": [model_input_video_neg], "audio": [model_input_audio], "t": timestep_input, **pos_args},
|
||||
{"vid": [model_input_video], "audio": [model_input_audio_neg], "t": timestep_input, **pos_args},
|
||||
]
|
||||
|
||||
pred_video_guided, pred_audio_guided = self.predict_noise_with_multi_branch_cfg(
|
||||
do_true_cfg=True,
|
||||
true_cfg_scale={
|
||||
"video_cfg_scale": self.video_cfg_scale,
|
||||
"video_ref_cfg_scale": self.video_ref_cfg_scale,
|
||||
"audio_cfg_scale": self.audio_cfg_scale,
|
||||
"audio_ref_cfg_scale": self.audio_ref_cfg_scale,
|
||||
},
|
||||
branches_kwargs=branches_kwargs,
|
||||
)
|
||||
video_noise = scheduler_video.step(
|
||||
pred_video_guided.unsqueeze(0), t_v, video_noise.unsqueeze(0), return_dict=False
|
||||
)[0].squeeze(0)
|
||||
audio_noise = scheduler_audio.step(
|
||||
pred_audio_guided.unsqueeze(0), t_a, audio_noise.unsqueeze(0), return_dict=False
|
||||
)[0].squeeze(0)
|
||||
return video_noise, audio_noise
|
||||
|
||||
def encode_prompt(
|
||||
self,
|
||||
prompt: str,
|
||||
video_negative_prompt: str = "",
|
||||
audio_negative_prompt: str = "",
|
||||
) -> tuple[torch.Tensor, torch.Tensor, torch.Tensor, torch.Tensor]:
|
||||
"""Encode the positive and negative prompts via ``self.text_model``.
|
||||
|
||||
This is the single text-encoder entrypoint so that the runner-level
|
||||
prompt-embedding cache (see
|
||||
``vllm_omni/diffusion/cache/prompt_embed_cache.py``) can transparently
|
||||
memoize results when the same prompts are submitted repeatedly (e.g.
|
||||
GRPO rollouts that sample the same prompt with different seeds).
|
||||
|
||||
Returns:
|
||||
(audio_pos, video_pos, video_neg, audio_neg) embeddings cast to
|
||||
``self.target_dtype``. ``audio_pos`` and ``video_pos`` are the
|
||||
same tensor, matching the original inline behavior.
|
||||
"""
|
||||
text_embeddings = self.text_model(
|
||||
[prompt, video_negative_prompt, audio_negative_prompt],
|
||||
device=self.device,
|
||||
)
|
||||
text_embeddings = [emb.to(self.target_dtype) for emb in text_embeddings]
|
||||
# Index 0 is the positive embedding, shared between video and audio branches.
|
||||
return text_embeddings[0], text_embeddings[0], text_embeddings[1], text_embeddings[2]
|
||||
|
||||
def forward(
|
||||
self,
|
||||
request: DiffusionRequestBatch,
|
||||
**kwargs,
|
||||
) -> DiffusionOutput:
|
||||
"""Main forward pass for DreamID-Omni pipeline for R2AV task."""
|
||||
# Extract parameters from request
|
||||
r_prompts = request.prompts[0]
|
||||
if isinstance(r_prompts, str):
|
||||
prompt = r_prompts
|
||||
video_negative_prompt = ""
|
||||
audio_negative_prompt = ""
|
||||
else:
|
||||
prompt = r_prompts.get("prompt", "")
|
||||
video_negative_prompt = r_prompts.get("video_negative_prompt", "")
|
||||
audio_negative_prompt = r_prompts.get("audio_negative_prompt", "")
|
||||
|
||||
multi_modal_data = r_prompts.get("multi_modal_data", {}) if not isinstance(r_prompts, str) else None
|
||||
raw_images = multi_modal_data.get("image", None) if multi_modal_data is not None else None
|
||||
raw_audios = multi_modal_data.get("audio", None) if multi_modal_data is not None else None
|
||||
if raw_images is None or raw_audios is None:
|
||||
raise ValueError("This model requires image and audio to run.")
|
||||
if not isinstance(raw_images, list):
|
||||
raw_images = [raw_images]
|
||||
if not isinstance(raw_audios, list):
|
||||
raw_audios = [raw_audios]
|
||||
if len(raw_images) != len(raw_audios):
|
||||
raise ValueError("'image' and 'audio' lists must have the same length.")
|
||||
|
||||
height = request.sampling_params.height
|
||||
width = request.sampling_params.width
|
||||
num_inference_steps = request.sampling_params.num_inference_steps
|
||||
shift = request.sampling_params.extra_args.get("shift", 5.0)
|
||||
solver_name = request.sampling_params.extra_args.get("solver_name", "unipc")
|
||||
seed = request.sampling_params.seed if request.sampling_params.seed is not None else 42
|
||||
|
||||
# 1. Prepare reference latents
|
||||
ref_vae_latents, ref_audio_lengths = self.load_image_latent_ref_ip_video(
|
||||
images=raw_images,
|
||||
audios=raw_audios,
|
||||
video_frame_height_width=(height, width),
|
||||
)
|
||||
|
||||
latents_ref_image = ref_vae_latents["image"]
|
||||
latents_ref_audio = ref_vae_latents["audio"]
|
||||
ref_ip_num = latents_ref_image.shape[1]
|
||||
ref_audio_length = latents_ref_audio.shape[0]
|
||||
|
||||
# 2. scheduler
|
||||
scheduler_video, timesteps_video = self.get_scheduler_time_steps(
|
||||
sampling_steps=num_inference_steps, device=self.device, solver_name=solver_name, shift=shift
|
||||
)
|
||||
scheduler_audio, timesteps_audio = self.get_scheduler_time_steps(
|
||||
sampling_steps=num_inference_steps, device=self.device, solver_name=solver_name, shift=shift
|
||||
)
|
||||
|
||||
# 3. text embedding
|
||||
text_embeddings_audio_pos, text_embeddings_video_pos, text_embeddings_video_neg, text_embeddings_audio_neg = (
|
||||
self.encode_prompt(
|
||||
prompt=prompt,
|
||||
video_negative_prompt=video_negative_prompt,
|
||||
audio_negative_prompt=audio_negative_prompt,
|
||||
)
|
||||
)
|
||||
|
||||
video_latent_h, video_latent_w = height // 16, width // 16
|
||||
|
||||
video_noise_len = self.video_latent_length + ref_ip_num
|
||||
audio_noise_len = self.audio_latent_length + ref_audio_length
|
||||
freqs_scaling_tensor = torch.tensor(
|
||||
self.video_latent_length / self.audio_latent_length, device=self.device, dtype=self.target_dtype
|
||||
)
|
||||
|
||||
video_noise = torch.randn(
|
||||
(self.video_latent_channel, video_noise_len, video_latent_h, video_latent_w),
|
||||
device=self.device,
|
||||
dtype=self.target_dtype,
|
||||
generator=torch.Generator(device=self.device).manual_seed(seed),
|
||||
)
|
||||
audio_noise = torch.randn(
|
||||
(audio_noise_len, self.audio_latent_channel),
|
||||
device=self.device,
|
||||
dtype=self.target_dtype,
|
||||
generator=torch.Generator(device=self.device).manual_seed(seed),
|
||||
)
|
||||
|
||||
_patch_size_h, _patch_size_w = self.model.video_model.patch_size[1], self.model.video_model.patch_size[2]
|
||||
|
||||
max_seq_len_video = (
|
||||
video_noise.shape[1] * video_noise.shape[2] * video_noise.shape[3] // (_patch_size_h * _patch_size_w)
|
||||
)
|
||||
|
||||
max_seq_len_audio = audio_noise_len
|
||||
|
||||
with torch.amp.autocast("cuda", enabled=self.target_dtype != torch.float32, dtype=self.target_dtype):
|
||||
video_noise, audio_noise = self.diffuse(
|
||||
video_noise,
|
||||
audio_noise,
|
||||
latents_ref_image,
|
||||
latents_ref_audio,
|
||||
timesteps_video,
|
||||
timesteps_audio,
|
||||
text_embeddings_video_pos,
|
||||
text_embeddings_video_neg,
|
||||
text_embeddings_audio_pos,
|
||||
text_embeddings_audio_neg,
|
||||
max_seq_len_video,
|
||||
max_seq_len_audio,
|
||||
freqs_scaling_tensor,
|
||||
ref_ip_num,
|
||||
ref_audio_length,
|
||||
ref_audio_lengths,
|
||||
scheduler_video,
|
||||
scheduler_audio,
|
||||
)
|
||||
|
||||
video_noise_for_decode = video_noise[:, :-ref_ip_num]
|
||||
audio_noise_for_decode = audio_noise[:-ref_audio_length, :]
|
||||
|
||||
audio_latents_for_vae = audio_noise_for_decode.unsqueeze(0).transpose(1, 2)
|
||||
generated_audio = self.vae_model_audio.wrapped_decode(audio_latents_for_vae).squeeze().cpu().float().numpy()
|
||||
|
||||
video_latents_for_vae = video_noise_for_decode.unsqueeze(0)
|
||||
generated_video = self.vae_model_video.wrapped_decode(video_latents_for_vae).squeeze(0).cpu().float().numpy()
|
||||
|
||||
return DiffusionOutput(output=(generated_video, generated_audio))
|
||||
@@ -1,694 +0,0 @@
|
||||
# Copyright 2024-2025 The Alibaba Wan Team Authors. All rights reserved.
|
||||
import math
|
||||
|
||||
import torch
|
||||
import torch.amp as amp
|
||||
import torch.nn as nn
|
||||
from diffusers.configuration_utils import ConfigMixin, register_to_config
|
||||
from diffusers.models.modeling_utils import ModelMixin
|
||||
from vllm.distributed import (
|
||||
get_tensor_model_parallel_world_size,
|
||||
tensor_model_parallel_all_reduce,
|
||||
)
|
||||
from vllm.model_executor.layers.linear import (
|
||||
ColumnParallelLinear,
|
||||
QKVParallelLinear,
|
||||
RowParallelLinear,
|
||||
)
|
||||
from vllm.model_executor.layers.quantization.base_config import QuantizationConfig
|
||||
|
||||
from vllm_omni.diffusion.attention.layer import Attention
|
||||
|
||||
try:
|
||||
from dreamid_omni.modules.model import (
|
||||
ChannelLastConv1d,
|
||||
ConvMLP,
|
||||
Head,
|
||||
MLPProj,
|
||||
ModulationAdd,
|
||||
WanLayerNorm,
|
||||
rope_apply, # diff from wan2.2
|
||||
rope_params,
|
||||
sinusoidal_embedding_1d,
|
||||
)
|
||||
except ImportError:
|
||||
raise ImportError("Failed to import from dependency 'dreamid_omni'.")
|
||||
|
||||
|
||||
class DistributedRMSNorm(nn.Module):
|
||||
"""RMSNorm that computes global RMS across tensor parallel ranks.
|
||||
|
||||
Mirrors `vllm_omni/diffusion/models/wan2_2/wan2_2_transformer.py::DistributedRMSNorm`
|
||||
"""
|
||||
|
||||
def __init__(self, hidden_size: int, eps: float = 1e-6):
|
||||
super().__init__()
|
||||
self.eps = eps
|
||||
self.weight = nn.Parameter(torch.ones(hidden_size))
|
||||
|
||||
def forward(self, x: torch.Tensor) -> torch.Tensor:
|
||||
tp_size = get_tensor_model_parallel_world_size()
|
||||
x_dtype = x.dtype
|
||||
x_float = x.float()
|
||||
local_sum_sq = x_float.pow(2).sum(dim=-1, keepdim=True)
|
||||
local_count = x.shape[-1]
|
||||
|
||||
if tp_size > 1:
|
||||
local_sum_sq = tensor_model_parallel_all_reduce(local_sum_sq)
|
||||
local_count = local_count * tp_size
|
||||
|
||||
rms = torch.sqrt(local_sum_sq / local_count + self.eps)
|
||||
return (x_float / rms * self.weight.float()).to(x_dtype)
|
||||
|
||||
|
||||
class WanSelfAttention(nn.Module):
|
||||
"""Optimized self-attention module using vLLM layers."""
|
||||
|
||||
def __init__(
|
||||
self,
|
||||
dim,
|
||||
num_heads,
|
||||
window_size=(-1, -1),
|
||||
qk_norm=True,
|
||||
eps=1e-6,
|
||||
*,
|
||||
quant_config: QuantizationConfig | None = None,
|
||||
prefix: str = "",
|
||||
):
|
||||
assert dim % num_heads == 0
|
||||
super().__init__()
|
||||
self.dim = dim
|
||||
self.head_dim = dim // num_heads
|
||||
self.window_size = window_size
|
||||
self.qk_norm = qk_norm
|
||||
self.eps = eps
|
||||
|
||||
self.to_qkv = QKVParallelLinear(
|
||||
hidden_size=dim,
|
||||
head_size=self.head_dim,
|
||||
total_num_heads=num_heads,
|
||||
bias=True,
|
||||
quant_config=quant_config,
|
||||
prefix=f"{prefix}.to_qkv" if prefix else "to_qkv",
|
||||
)
|
||||
self.num_heads = self.to_qkv.num_heads
|
||||
self.num_kv_heads = self.to_qkv.num_kv_heads
|
||||
self.tp_inner_dim = self.num_heads * self.head_dim
|
||||
self.tp_kv_dim = self.num_kv_heads * self.head_dim
|
||||
|
||||
self.o = RowParallelLinear(
|
||||
dim,
|
||||
dim,
|
||||
bias=True,
|
||||
input_is_parallel=True,
|
||||
return_bias=False,
|
||||
quant_config=quant_config,
|
||||
prefix=f"{prefix}.o" if prefix else "o",
|
||||
)
|
||||
self.norm_q = DistributedRMSNorm(self.tp_inner_dim, eps=eps) if qk_norm else nn.Identity()
|
||||
self.norm_k = DistributedRMSNorm(self.tp_kv_dim, eps=eps) if qk_norm else nn.Identity()
|
||||
# Unified attention layer
|
||||
self.attn = Attention(
|
||||
num_heads=self.num_heads,
|
||||
head_size=self.head_dim,
|
||||
num_kv_heads=self.num_kv_heads,
|
||||
softmax_scale=1.0 / (self.head_dim**0.5),
|
||||
causal=False,
|
||||
prefix=prefix,
|
||||
)
|
||||
|
||||
# query, key, value function
|
||||
def qkv_fn(self, x):
|
||||
b, s = x.shape[:2]
|
||||
qkv, _ = self.to_qkv(x)
|
||||
q_size = self.num_heads * self.head_dim
|
||||
kv_size = self.num_kv_heads * self.head_dim
|
||||
q, k, v = qkv.split([q_size, kv_size, kv_size], dim=-1)
|
||||
|
||||
q = self.norm_q(q).view(b, s, self.num_heads, self.head_dim)
|
||||
k = self.norm_k(k).view(b, s, self.num_kv_heads, self.head_dim)
|
||||
v = v.view(b, s, self.num_kv_heads, self.head_dim)
|
||||
return q, k, v
|
||||
|
||||
def forward(self, x, seq_lens, grid_sizes, freqs, ref_lengths=None, freqs_scaling=None):
|
||||
r"""
|
||||
Args:
|
||||
x(Tensor): Shape [B, L, C]
|
||||
seq_lens(Tensor): Shape [B]
|
||||
grid_sizes(Tensor): Shape [B, 3], the second dimension contains (F, H, W)
|
||||
freqs(Tensor): Rope freqs, shape [1024, C / num_heads / 2]
|
||||
ref_lengths(Tensor): Shape [B]
|
||||
"""
|
||||
q, k, v = self.qkv_fn(x)
|
||||
|
||||
x = self.attn(
|
||||
rope_apply(q, grid_sizes, freqs, ref_lengths=ref_lengths, freqs_scaling=freqs_scaling),
|
||||
rope_apply(k, grid_sizes, freqs, ref_lengths=ref_lengths, freqs_scaling=freqs_scaling),
|
||||
v,
|
||||
)
|
||||
|
||||
# output
|
||||
x = x.flatten(2)
|
||||
x = self.o(x)
|
||||
return x
|
||||
|
||||
|
||||
class WanT2VCrossAttention(nn.Module):
|
||||
"""Text -> latent cross-attention."""
|
||||
|
||||
def __init__(
|
||||
self,
|
||||
dim,
|
||||
num_heads,
|
||||
window_size=(-1, -1),
|
||||
qk_norm=True,
|
||||
eps=1e-6,
|
||||
*,
|
||||
quant_config: QuantizationConfig | None = None,
|
||||
prefix: str = "",
|
||||
):
|
||||
super().__init__()
|
||||
assert dim % num_heads == 0
|
||||
self.dim = dim
|
||||
self.head_dim = dim // num_heads
|
||||
self.window_size = window_size
|
||||
self.qk_norm = qk_norm
|
||||
self.eps = eps
|
||||
|
||||
# Separate Q/K/V — Q's input differs from K/V's input (text context)
|
||||
self.q = ColumnParallelLinear(
|
||||
dim,
|
||||
dim,
|
||||
bias=True,
|
||||
gather_output=False,
|
||||
return_bias=False,
|
||||
quant_config=quant_config,
|
||||
prefix=f"{prefix}.q" if prefix else "q",
|
||||
)
|
||||
self.k = ColumnParallelLinear(
|
||||
dim,
|
||||
dim,
|
||||
bias=True,
|
||||
gather_output=False,
|
||||
return_bias=False,
|
||||
quant_config=quant_config,
|
||||
prefix=f"{prefix}.k" if prefix else "k",
|
||||
)
|
||||
self.v = ColumnParallelLinear(
|
||||
dim,
|
||||
dim,
|
||||
bias=True,
|
||||
gather_output=False,
|
||||
return_bias=False,
|
||||
quant_config=quant_config,
|
||||
prefix=f"{prefix}.v" if prefix else "v",
|
||||
)
|
||||
|
||||
tp_size = get_tensor_model_parallel_world_size()
|
||||
self.num_heads = num_heads // tp_size
|
||||
self.num_kv_heads = self.num_heads
|
||||
self.tp_inner_dim = self.num_heads * self.head_dim
|
||||
|
||||
self.norm_q = DistributedRMSNorm(self.tp_inner_dim, eps=eps) if qk_norm else nn.Identity()
|
||||
self.norm_k = DistributedRMSNorm(self.tp_inner_dim, eps=eps) if qk_norm else nn.Identity()
|
||||
|
||||
self.o = RowParallelLinear(
|
||||
dim,
|
||||
dim,
|
||||
bias=True,
|
||||
input_is_parallel=True,
|
||||
return_bias=False,
|
||||
quant_config=quant_config,
|
||||
prefix=f"{prefix}.o" if prefix else "o",
|
||||
)
|
||||
|
||||
# Kept for interface parity with WanSelfAttention;
|
||||
# the active cross-attention path in DreamID-Omni uses the shared FusionModel.attn.
|
||||
# Disable kv quant to prevent affecting quality.
|
||||
self.attn = Attention(
|
||||
num_heads=self.num_heads,
|
||||
head_size=self.head_dim,
|
||||
num_kv_heads=self.num_kv_heads,
|
||||
softmax_scale=1.0 / (self.head_dim**0.5),
|
||||
causal=False,
|
||||
prefix=prefix,
|
||||
disable_kv_quant=True,
|
||||
)
|
||||
|
||||
def qkv_fn(self, x, context):
|
||||
b = x.size(0)
|
||||
|
||||
q = self.q(x)
|
||||
k = self.k(context)
|
||||
v = self.v(context)
|
||||
q = self.norm_q(q).view(b, -1, self.num_heads, self.head_dim)
|
||||
k = self.norm_k(k).view(b, -1, self.num_kv_heads, self.head_dim)
|
||||
v = v.view(b, -1, self.num_kv_heads, self.head_dim)
|
||||
|
||||
return q, k, v
|
||||
|
||||
def forward(self, x, context, context_lens):
|
||||
r"""
|
||||
Args:
|
||||
x(Tensor): Shape [B, L1, C]
|
||||
context(Tensor): Shape [B, L2, C]
|
||||
context_lens(Tensor): Shape [B]
|
||||
"""
|
||||
q, k, v = self.qkv_fn(x, context)
|
||||
|
||||
# compute attention
|
||||
x = self.attn(q, k, v)
|
||||
|
||||
# output
|
||||
x = x.flatten(2)
|
||||
x = self.o(x)
|
||||
return x
|
||||
|
||||
|
||||
class WanI2VCrossAttention(nn.Module):
|
||||
def __init__(self, *args, **kwargs):
|
||||
super().__init__()
|
||||
|
||||
def forward(self, *args, **kwargs):
|
||||
raise RuntimeError("WanT2VCrossAttention is currently disabled and should not be used.")
|
||||
|
||||
|
||||
WAN_CROSSATTENTION_CLASSES = {
|
||||
"t2v_cross_attn": WanT2VCrossAttention,
|
||||
"i2v_cross_attn": WanI2VCrossAttention,
|
||||
}
|
||||
|
||||
|
||||
class WanAttentionBlock(nn.Module):
|
||||
def __init__(
|
||||
self,
|
||||
cross_attn_type,
|
||||
dim,
|
||||
ffn_dim,
|
||||
num_heads,
|
||||
window_size=(-1, -1),
|
||||
qk_norm=True,
|
||||
cross_attn_norm=False,
|
||||
eps=1e-6,
|
||||
additional_emb_length=None,
|
||||
*,
|
||||
quant_config: QuantizationConfig | None = None,
|
||||
prefix: str = "",
|
||||
):
|
||||
super().__init__()
|
||||
self.dim = dim
|
||||
self.ffn_dim = ffn_dim
|
||||
self.num_heads = num_heads
|
||||
self.window_size = window_size
|
||||
self.qk_norm = qk_norm
|
||||
self.cross_attn_norm = cross_attn_norm
|
||||
self.eps = eps
|
||||
|
||||
# layers
|
||||
self.norm1 = WanLayerNorm(dim, eps)
|
||||
self.self_attn = WanSelfAttention(
|
||||
dim,
|
||||
num_heads,
|
||||
window_size,
|
||||
qk_norm,
|
||||
eps,
|
||||
quant_config=quant_config,
|
||||
prefix=f"{prefix}.self_attn" if prefix else "self_attn",
|
||||
)
|
||||
self.norm3 = WanLayerNorm(dim, eps, elementwise_affine=True) if cross_attn_norm else nn.Identity()
|
||||
if cross_attn_type == "i2v_cross_attn":
|
||||
assert additional_emb_length is not None, "additional_emb_length should be specified for i2v_cross_attn"
|
||||
self.cross_attn = WanI2VCrossAttention(dim, num_heads, (-1, -1), qk_norm, eps, additional_emb_length)
|
||||
else:
|
||||
assert additional_emb_length is None, "additional_emb_length should be None for t2v_cross_attn"
|
||||
self.cross_attn = WanT2VCrossAttention(
|
||||
dim,
|
||||
num_heads,
|
||||
(-1, -1),
|
||||
qk_norm,
|
||||
eps,
|
||||
quant_config=quant_config,
|
||||
prefix=f"{prefix}.cross_attn" if prefix else "cross_attn",
|
||||
)
|
||||
self.norm2 = WanLayerNorm(dim, eps)
|
||||
# FFN layout: Stay compatible with upstream DreamID-Omni checkpoints
|
||||
self.ffn = nn.Sequential(
|
||||
ColumnParallelLinear(
|
||||
dim,
|
||||
ffn_dim,
|
||||
bias=True,
|
||||
gather_output=False,
|
||||
return_bias=False,
|
||||
quant_config=quant_config,
|
||||
prefix=f"{prefix}.ffn.0" if prefix else "ffn.0",
|
||||
),
|
||||
nn.GELU(approximate="tanh"),
|
||||
RowParallelLinear(
|
||||
ffn_dim,
|
||||
dim,
|
||||
bias=True,
|
||||
input_is_parallel=True,
|
||||
return_bias=False,
|
||||
quant_config=quant_config,
|
||||
prefix=f"{prefix}.ffn.2" if prefix else "ffn.2",
|
||||
),
|
||||
)
|
||||
|
||||
# modulation
|
||||
self.modulation = ModulationAdd(dim, 6)
|
||||
|
||||
def forward(
|
||||
self,
|
||||
x,
|
||||
e,
|
||||
seq_lens,
|
||||
grid_sizes,
|
||||
freqs,
|
||||
context,
|
||||
context_lens,
|
||||
):
|
||||
r"""
|
||||
Args:
|
||||
x(Tensor): Shape [B, L, C]
|
||||
e(Tensor): Shape [B, L1, 6, C]
|
||||
seq_lens(Tensor): Shape [B], length of each sequence in batch
|
||||
grid_sizes(Tensor): Shape [B, 3], the second dimension contains (F, H, W)
|
||||
freqs(Tensor): Rope freqs, shape [1024, C / num_heads / 2]
|
||||
"""
|
||||
assert e.dtype == torch.bfloat16
|
||||
assert len(e.shape) == 4 and e.size(2) == 6 and e.shape[1] == x.shape[1], f"{e.shape}, {x.shape}"
|
||||
with amp.autocast("cuda", dtype=torch.bfloat16):
|
||||
e = self.modulation(e).chunk(6, dim=2)
|
||||
assert e[0].dtype == torch.bfloat16
|
||||
|
||||
# self-attention
|
||||
y = self.self_attn(
|
||||
self.norm1(x).bfloat16() * (1 + e[1].squeeze(2)) + e[0].squeeze(2), seq_lens, grid_sizes, freqs
|
||||
)
|
||||
with amp.autocast("cuda", dtype=torch.bfloat16):
|
||||
x = x + y * e[2].squeeze(2)
|
||||
|
||||
# cross-attention & ffn function
|
||||
def cross_attn_ffn(x, context, context_lens, e):
|
||||
x = x + self.cross_attn(self.norm3(x), context, context_lens)
|
||||
y = self.ffn(self.norm2(x).bfloat16() * (1 + e[4].squeeze(2)) + e[3].squeeze(2))
|
||||
with amp.autocast("cuda", dtype=torch.bfloat16):
|
||||
x = x + y * e[5].squeeze(2)
|
||||
return x
|
||||
|
||||
x = cross_attn_ffn(x, context, context_lens, e)
|
||||
return x
|
||||
|
||||
|
||||
class WanModel(ModelMixin, ConfigMixin):
|
||||
r"""
|
||||
Wan diffusion backbone supporting both text-to-video and image-to-video, text-to-audio.
|
||||
"""
|
||||
|
||||
@register_to_config
|
||||
def __init__(
|
||||
self,
|
||||
model_type="t2v",
|
||||
patch_size=(1, 2, 2),
|
||||
text_len=512,
|
||||
in_dim=16,
|
||||
dim=2048,
|
||||
ffn_dim=8192,
|
||||
freq_dim=256,
|
||||
text_dim=4096,
|
||||
additional_emb_dim=None,
|
||||
additional_emb_length=None,
|
||||
out_dim=16,
|
||||
num_heads=16,
|
||||
num_layers=32,
|
||||
window_size=(-1, -1),
|
||||
qk_norm=True,
|
||||
cross_attn_norm=True,
|
||||
temporal_rope_scaling_factor=1.0,
|
||||
eps=1e-6,
|
||||
*,
|
||||
quant_config: QuantizationConfig | None = None,
|
||||
prefix: str = "",
|
||||
):
|
||||
r"""
|
||||
Initialize the diffusion model backbone.
|
||||
|
||||
Args:
|
||||
model_type (`str`, *optional*, defaults to 't2v'):
|
||||
Model variant - 't2v' (text-to-video) or 'i2v' (image-to-video)
|
||||
patch_size (`tuple`, *optional*, defaults to (1, 2, 2)):
|
||||
3D patch dimensions for video embedding (t_patch, h_patch, w_patch)
|
||||
text_len (`int`, *optional*, defaults to 512):
|
||||
Fixed length for text embeddings
|
||||
in_dim (`int`, *optional*, defaults to 16):
|
||||
Input video channels (C_in)
|
||||
dim (`int`, *optional*, defaults to 2048):
|
||||
Hidden dimension of the transformer
|
||||
ffn_dim (`int`, *optional*, defaults to 8192):
|
||||
Intermediate dimension in feed-forward network
|
||||
freq_dim (`int`, *optional*, defaults to 256):
|
||||
Dimension for sinusoidal time embeddings
|
||||
text_dim (`int`, *optional*, defaults to 4096):
|
||||
Input dimension for text embeddings
|
||||
out_dim (`int`, *optional*, defaults to 16):
|
||||
Output video channels (C_out)
|
||||
num_heads (`int`, *optional*, defaults to 16):
|
||||
Number of attention heads
|
||||
num_layers (`int`, *optional*, defaults to 32):
|
||||
Number of transformer blocks
|
||||
window_size (`tuple`, *optional*, defaults to (-1, -1)):
|
||||
Window size for local attention (-1 indicates global attention)
|
||||
qk_norm (`bool`, *optional*, defaults to True):
|
||||
Enable query/key normalization
|
||||
cross_attn_norm (`bool`, *optional*, defaults to False):
|
||||
Enable cross-attention normalization
|
||||
eps (`float`, *optional*, defaults to 1e-6):
|
||||
Epsilon value for normalization layers
|
||||
"""
|
||||
|
||||
super().__init__()
|
||||
|
||||
assert model_type in [
|
||||
"t2v",
|
||||
"i2v",
|
||||
"t2a",
|
||||
"tt2a",
|
||||
"ti2v",
|
||||
] ## tt2a means text transcript + text description to audio (to support both TTS and T2A
|
||||
self.model_type = model_type
|
||||
is_audio_type = "a" in self.model_type
|
||||
is_video_type = "v" in self.model_type
|
||||
assert is_audio_type ^ is_video_type, "Either audio or video model should be specified"
|
||||
if is_audio_type:
|
||||
## audio model
|
||||
assert len(patch_size) == 1 and patch_size[0] == 1, (
|
||||
"Audio model should only accept 1 dimensional input, and we dont do patchify"
|
||||
)
|
||||
|
||||
self.patch_size = patch_size
|
||||
self.text_len = text_len
|
||||
self.in_dim = in_dim
|
||||
self.dim = dim
|
||||
self.ffn_dim = ffn_dim
|
||||
self.freq_dim = freq_dim
|
||||
self.text_dim = text_dim
|
||||
self.out_dim = out_dim
|
||||
self.num_heads = num_heads
|
||||
self.num_layers = num_layers
|
||||
self.window_size = window_size
|
||||
self.qk_norm = qk_norm
|
||||
self.cross_attn_norm = cross_attn_norm
|
||||
self.eps = eps
|
||||
self.temporal_rope_scaling_factor = temporal_rope_scaling_factor
|
||||
self.is_audio_type = is_audio_type
|
||||
self.is_video_type = is_video_type
|
||||
# embeddings
|
||||
if is_audio_type:
|
||||
## hardcoded to MMAudio
|
||||
self.patch_embedding = nn.Sequential(
|
||||
ChannelLastConv1d(in_dim, dim, kernel_size=7, padding=3),
|
||||
nn.SiLU(),
|
||||
ConvMLP(dim, dim * 4, kernel_size=7, padding=3),
|
||||
)
|
||||
else:
|
||||
self.patch_embedding = nn.Conv3d(in_dim, dim, kernel_size=patch_size, stride=patch_size)
|
||||
|
||||
self.text_embedding = nn.Sequential(nn.Linear(text_dim, dim), nn.GELU(approximate="tanh"), nn.Linear(dim, dim))
|
||||
|
||||
self.time_embedding = nn.Sequential(nn.Linear(freq_dim, dim), nn.SiLU(), nn.Linear(dim, dim))
|
||||
self.time_projection = nn.Sequential(nn.SiLU(), nn.Linear(dim, dim * 6))
|
||||
|
||||
# blocks
|
||||
## so i2v and tt2a share the same cross attention while t2v and t2a share the same cross attention
|
||||
cross_attn_type = "t2v_cross_attn" if model_type in ["t2v", "t2a", "ti2v"] else "i2v_cross_attn"
|
||||
|
||||
if cross_attn_type == "t2v_cross_attn":
|
||||
assert additional_emb_dim is None and additional_emb_length is None, (
|
||||
"additional_emb_length should be None for t2v and t2a model"
|
||||
)
|
||||
else:
|
||||
assert additional_emb_dim is not None and additional_emb_length is not None, (
|
||||
"additional_emb_length should be specified for i2v and tt2a model"
|
||||
)
|
||||
|
||||
self.blocks = nn.ModuleList(
|
||||
[
|
||||
WanAttentionBlock(
|
||||
cross_attn_type,
|
||||
dim,
|
||||
ffn_dim,
|
||||
num_heads,
|
||||
window_size,
|
||||
qk_norm,
|
||||
cross_attn_norm,
|
||||
eps,
|
||||
additional_emb_length,
|
||||
quant_config=quant_config,
|
||||
prefix=f"{prefix}.blocks.{layer_idx}" if prefix else f"blocks.{layer_idx}",
|
||||
)
|
||||
for layer_idx in range(num_layers)
|
||||
]
|
||||
)
|
||||
|
||||
# head
|
||||
self.head = Head(dim, out_dim, patch_size, eps)
|
||||
|
||||
self.set_rope_params()
|
||||
|
||||
if model_type in ["i2v", "tt2a"]:
|
||||
self.img_emb = MLPProj(additional_emb_dim, dim)
|
||||
|
||||
def set_rope_params(self):
|
||||
# buffers (don't use register_buffer otherwise dtype will be changed in to())
|
||||
dim = self.dim
|
||||
num_heads = self.num_heads
|
||||
assert (dim % num_heads) == 0 and (dim // num_heads) % 2 == 0
|
||||
d = dim // num_heads
|
||||
|
||||
if self.is_audio_type:
|
||||
## to be determined
|
||||
# self.freqs = rope_params(1024, d, freqs_scaling=temporal_rope_scaling_factor)
|
||||
self.freqs = rope_params(1024, d - 4 * (d // 6), freqs_scaling=self.temporal_rope_scaling_factor)
|
||||
else:
|
||||
self.freqs = torch.cat(
|
||||
[rope_params(1024, d - 4 * (d // 6)), rope_params(1024, 2 * (d // 6)), rope_params(1024, 2 * (d // 6))],
|
||||
dim=1,
|
||||
)
|
||||
|
||||
def prepare_transformer_block_kwargs(
|
||||
self,
|
||||
x,
|
||||
t,
|
||||
context,
|
||||
seq_len,
|
||||
ref_lengths=None,
|
||||
freqs_scaling=None,
|
||||
):
|
||||
# params
|
||||
## need to change!
|
||||
device = next(self.patch_embedding.parameters()).device
|
||||
|
||||
if self.is_audio_type and freqs_scaling is not None:
|
||||
if isinstance(freqs_scaling, torch.Tensor):
|
||||
scale_val = freqs_scaling.item()
|
||||
else:
|
||||
scale_val = freqs_scaling
|
||||
d = self.dim // self.num_heads
|
||||
|
||||
current_freqs = rope_params(1024, d - 4 * (d // 6), freqs_scaling=scale_val).to(device)
|
||||
else:
|
||||
current_freqs = self.freqs
|
||||
if current_freqs.device != device:
|
||||
current_freqs = current_freqs.to(device)
|
||||
|
||||
# embeddings
|
||||
x = [self.patch_embedding(u.unsqueeze(0)) for u in x] ## x is list of [B L D] or [B C F H W]
|
||||
if self.is_audio_type:
|
||||
# [B, 1]
|
||||
grid_sizes = torch.stack([torch.tensor(u.shape[1:2], dtype=torch.long) for u in x])
|
||||
else:
|
||||
# [B, 3]
|
||||
grid_sizes = torch.stack([torch.tensor(u.shape[2:], dtype=torch.long) for u in x])
|
||||
x = [u.flatten(2).transpose(1, 2) for u in x] # [B C F H W] -> [B (F H W) C] -> [B L C]
|
||||
|
||||
seq_lens = torch.tensor([u.size(1) for u in x], dtype=torch.long)
|
||||
assert seq_lens.max() <= seq_len, f"Sequence length {seq_lens.max()} exceeds maximum {seq_len}."
|
||||
x = torch.cat(
|
||||
[torch.cat([u, u.new_zeros(1, seq_len - u.size(1), u.size(2))], dim=1) for u in x]
|
||||
) # single [B, L, C]
|
||||
|
||||
# time embeddings
|
||||
if t.dim() == 1:
|
||||
t = t.unsqueeze(1).expand(t.size(0), seq_len)
|
||||
with amp.autocast("cuda", dtype=torch.bfloat16):
|
||||
bt = t.size(0)
|
||||
t = t.flatten()
|
||||
e = self.time_embedding(sinusoidal_embedding_1d(self.freq_dim, t).unflatten(0, (bt, seq_len)).bfloat16())
|
||||
e0 = self.time_projection(e).unflatten(2, (6, self.dim)) # [1, 26784, 6, 3072] - B, seq_len, 6, dim
|
||||
assert e.dtype == torch.bfloat16 and e0.dtype == torch.bfloat16
|
||||
|
||||
# context
|
||||
context_lens = None
|
||||
context = self.text_embedding(
|
||||
torch.stack([torch.cat([u, u.new_zeros(self.text_len - u.size(0), u.size(1))]) for u in context])
|
||||
)
|
||||
|
||||
# arguments
|
||||
kwargs = dict(
|
||||
e=e0,
|
||||
seq_lens=seq_lens,
|
||||
grid_sizes=grid_sizes,
|
||||
freqs=current_freqs,
|
||||
context=context,
|
||||
context_lens=context_lens,
|
||||
ref_lengths=ref_lengths,
|
||||
freqs_scaling=freqs_scaling,
|
||||
)
|
||||
|
||||
return x, e, kwargs
|
||||
|
||||
def post_transformer_block_out(self, x, grid_sizes, e):
|
||||
# head
|
||||
x = self.head(x, e)
|
||||
# unpatchify
|
||||
if self.is_audio_type:
|
||||
## grid_sizes is [B 1] where 1 is L,
|
||||
# converting grid_sizes from [B 1] -> [B]
|
||||
grid_sizes = [gs[0] for gs in grid_sizes]
|
||||
assert len(x) == len(grid_sizes)
|
||||
x = [u[:gs] for u, gs in zip(x, grid_sizes)]
|
||||
else:
|
||||
## grid_sizes is [B 3] where 3 is F H w
|
||||
x = self.unpatchify(x, grid_sizes)
|
||||
|
||||
return [u.bfloat16() for u in x]
|
||||
|
||||
def forward(self, *args, **kwargs):
|
||||
raise NotImplementedError("WanModel model does not support forward pass.")
|
||||
|
||||
def unpatchify(self, x, grid_sizes) -> list[torch.Tensor]:
|
||||
r"""
|
||||
Reconstruct video tensors from patch embeddings.
|
||||
|
||||
Args:
|
||||
x (List[Tensor]):
|
||||
List of patchified features, each with shape [L, C_out * prod(patch_size)]
|
||||
grid_sizes (Tensor):
|
||||
Original spatial-temporal grid dimensions before patching,
|
||||
shape [B, 3] (3 dimensions correspond to F_patches, H_patches, W_patches)
|
||||
|
||||
Returns:
|
||||
List[Tensor]:
|
||||
Reconstructed video tensors with shape [C_out, F, H / 8, W / 8]
|
||||
"""
|
||||
|
||||
c = self.out_dim
|
||||
out = []
|
||||
for u, v in zip(x, grid_sizes.tolist()):
|
||||
# v is [F H w] F * H * 80, 100, it was right padded by 20.
|
||||
u = u[: math.prod(v)].view(*v, *self.patch_size, c)
|
||||
u = torch.einsum("fhwpqrc->cfphqwr", u)
|
||||
u = u.reshape(c, *[i * j for i, j in zip(v, self.patch_size)])
|
||||
out.append(u)
|
||||
# out is list of [C F H W]
|
||||
return out
|
||||
@@ -1 +0,0 @@
|
||||
# SPDX-License-Identifier: Apache-2.0
|
||||
File diff suppressed because it is too large.
Load diff
File diff suppressed because it is too large.
Load diff
Reference in new issue
Block a user