anodyine/vllm-omni
1
0 Fork 0
Code Issues 4 Pull Requests 2 Actions Packages Projects Releases Wiki Activity
Files
40dcfdbb4343fca917b34287cba620eb436a856d
vllm-omni/vllm_omni/diffusion/models/dreamid_omni/pipeline_dreamid_omni.py
T

609 lines
24 KiB
Python
Raw Blame History

# 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))
Reference in new issue View Git Blame Copy Permalink