anodyine/vllm-omni
1
0 Fork 0
Code Issues 4 Pull Requests 2 Actions Packages Projects Releases Wiki Activity

Merge branch 'main' into lance-model-support

Browse Source
This commit is contained in:
SYLAR authored and GitHub committed 2026-05-22 04:13:58 +08:00
2 parents 7b885f9d13 e949ccf087
commit f0cec3ecd4
224 files changed
+19202 -7026
Download Patch File Download Diff File Expand all files Collapse all files

No files matched your search

vllm_omni/diffusion/models/flux2_klein/pipeline_flux2_klein.py
+9 -1
View File
@@ -701,6 +701,10 @@ class Flux2KleinPipeline(nn.Module, CFGParallelMixin, SupportImageInput, Diffusi
raise ValueError(
"Provide either `prompt` or `prompt_embeds`. Cannot leave both `prompt` and `prompt_embeds` undefined."
)
elif isinstance(prompt, str) and not prompt.strip():
raise ValueError("`prompt` cannot be empty or whitespace-only.")
elif isinstance(prompt, list) and any(isinstance(p, str) and not p.strip() for p in prompt):
raise ValueError("`prompt` cannot contain empty or whitespace-only strings.")
elif prompt is not None and (not isinstance(prompt, str) and not isinstance(prompt, list)):
raise ValueError(f"`prompt` has to be of type `str` or `list` but is {type(prompt)}")
@@ -861,7 +865,11 @@ class Flux2KleinPipeline(nn.Module, CFGParallelMixin, SupportImageInput, Diffusi
height = req.sampling_params.height or height
width = req.sampling_params.width or width
num_inference_steps = req.sampling_params.num_inference_steps or num_inference_steps
num_inference_steps = (
req.sampling_params.num_inference_steps
if req.sampling_params.num_inference_steps is not None
else num_inference_steps
)
sigmas = req.sampling_params.sigmas or sigmas
guidance_scale = (
req.sampling_params.guidance_scale if req.sampling_params.guidance_scale is not None else guidance_scale
vllm_omni/diffusion/models/hunyuan_image3/hunyuan_image3_transformer.py
+15 -1
View File
@@ -72,6 +72,7 @@ from vllm_omni.diffusion.distributed.sp_plan import (
SequenceParallelOutput,
)
from vllm_omni.diffusion.distributed.utils import get_local_device
from vllm_omni.diffusion.forward_context import set_forward_context_denoise_step_idx
from vllm_omni.diffusion.layers.rope import RotaryEmbedding
from vllm_omni.diffusion.models.hunyuan_image3.hunyuan_fused_moe import HunyuanFusedMoE
from vllm_omni.model_executor.layers.timestep_embedding import timestep_embedding
@@ -868,7 +869,15 @@ class ImageKVCacheManager:
Manages specialized caching and updating of KV-Cache for image tokens in multimodal models.
"""
def __init__(self, num_heads: int, num_kv_heads: int, head_dim: int, scaling: float, image_token_len: int = 4097):
def __init__(
self,
num_heads: int,
num_kv_heads: int,
head_dim: int,
scaling: float,
image_token_len: int = 4097,
prefix: str = "",
):
"""
Args:
image_token_len: Number of tokens per image (including special placeholders),
@@ -892,6 +901,7 @@ class ImageKVCacheManager:
causal=False,
softmax_scale=self.scaling,
num_kv_heads=self.num_kv_heads,
prefix=f"{prefix}.attn" if prefix else "",
)
def _cache_prompt_kv(
@@ -1650,6 +1660,7 @@ class HunYuanAttention(nn.Module):
num_kv_heads=self.num_kv_heads,
scaling=self.scaling,
image_token_len=4097,
prefix=f"{prefix}.image_attn",
)
self.image_rope2d_emb = HunYuanRotary2DEmbedder(
num_heads=self.num_heads,
@@ -3021,6 +3032,7 @@ class HunyuanImage3Text2ImagePipeline(DiffusionPipeline):
with self.progress_bar(total=num_inference_steps) as progress_bar:
for i, t in enumerate(timesteps):
set_forward_context_denoise_step_idx(i)
if cfg_parallel_ready:
# CFG parallel: each rank forwards its own branch (no batch doubling)
latent_model_input = latents
@@ -3109,6 +3121,8 @@ class HunyuanImage3Text2ImagePipeline(DiffusionPipeline):
if i == len(timesteps) - 1 or ((i + 1) > num_warmup_steps and (i + 1) % self.scheduler.order == 0):
progress_bar.update()
set_forward_context_denoise_step_idx(None)
if hasattr(self.vae.config, "scaling_factor") and self.vae.config.scaling_factor:
latents = latents / self.vae.config.scaling_factor
if hasattr(self.vae.config, "shift_factor") and self.vae.config.shift_factor:
vllm_omni/diffusion/models/hunyuan_image3/pipeline_hunyuan_image3.py
+7 -2
View File
@@ -30,7 +30,6 @@ from vllm_omni.diffusion.request import OmniDiffusionRequest
from vllm_omni.inputs.data import OmniTextPrompt
from vllm_omni.model_executor.models.hunyuan_image3.siglip2 import Siglip2VisionTransformer
from .autoencoder import AutoencoderKLConv3D
from .hunyuan_image3_tokenizer import TokenizerWrapper
from .hunyuan_image3_transformer import (
CausalMMOutputWithPast,
@@ -343,7 +342,13 @@ class HunyuanImage3Pipeline(
quant_config = od_config.quantization_config
self.model = HunyuanImage3Model(self.hf_config, quant_config=quant_config)
self.transformer = self.model
self.vae = AutoencoderKLConv3D.from_config(self.hf_config.vae)
# Lazy import to break circular dependency:
# autoencoder_kl_hunyuan -> hunyuan_image3/__init__ -> pipeline_hunyuan_image3 -> autoencoder_kl_hunyuan
from vllm_omni.diffusion.distributed.autoencoders.autoencoder_kl_hunyuan import ( # noqa: PLC0415
DistributedAutoencoderKLHunyuan,
)
self.vae = DistributedAutoencoderKLHunyuan.from_config(self.hf_config.vae)
self.vae.use_spatial_tiling = self.od_config.vae_use_tiling
self._pipeline = None
self._tkwrapper = TokenizerWrapper(od_config.model)
vllm_omni/diffusion/models/ltx2/ltx2_transformer.py
+81 -6
View File
@@ -17,7 +17,7 @@ import inspect
from collections.abc import Iterable
from dataclasses import dataclass
from types import SimpleNamespace
from typing import Any
from typing import TYPE_CHECKING, Any
import torch
import torch.distributed
@@ -39,6 +39,9 @@ from vllm.model_executor.layers.linear import (
)
from vllm.model_executor.model_loader.weight_utils import default_weight_loader
if TYPE_CHECKING:
from vllm.model_executor.layers.quantization.base_config import QuantizationConfig
from vllm_omni.diffusion.attention.backends.abstract import AttentionMetadata
from vllm_omni.diffusion.attention.layer import Attention
from vllm_omni.diffusion.distributed.hsdp_utils import is_transformer_block_module
@@ -173,7 +176,16 @@ class LTX2AdaLayerNormSingle(nn.Module):
class ColumnParallelApproxGELU(nn.Module):
def __init__(self, dim_in: int, dim_out: int, *, approximate: str, bias: bool = True):
def __init__(
self,
dim_in: int,
dim_out: int,
*,
approximate: str,
bias: bool = True,
quant_config: "QuantizationConfig | None" = None,
prefix: str = "",
):
super().__init__()
self.proj = ColumnParallelLinear(
dim_in,
@@ -181,6 +193,8 @@ class ColumnParallelApproxGELU(nn.Module):
bias=bias,
gather_output=False,
return_bias=False,
quant_config=quant_config,
prefix=f"{prefix}.proj" if prefix else "proj",
)
self.approximate = approximate
@@ -200,6 +214,8 @@ class LTX2FeedForward(nn.Module):
bias: bool = True,
dropout: float = 0.0,
final_dropout: bool = False,
quant_config: "QuantizationConfig | None" = None,
prefix: str = "",
) -> None:
super().__init__()
@@ -211,13 +227,22 @@ class LTX2FeedForward(nn.Module):
dropout_layer: nn.Module = nn.Dropout(dropout) if dropout > 0 else nn.Identity()
layers: list[nn.Module] = [
ColumnParallelApproxGELU(dim, inner_dim, approximate="tanh", bias=bias),
ColumnParallelApproxGELU(
dim,
inner_dim,
approximate="tanh",
bias=bias,
quant_config=quant_config,
prefix=f"{prefix}.net.0" if prefix else "net.0",
),
dropout_layer,
RowParallelLinear(
inner_dim,
dim_out,
input_is_parallel=True,
return_bias=False,
quant_config=quant_config,
prefix=f"{prefix}.net.2" if prefix else "net.2",
),
]
if final_dropout:
@@ -515,6 +540,9 @@ class LTX2Attention(torch.nn.Module):
rope_type: str = "interleaved",
apply_gated_attention: bool = False,
processor=None,
quant_config: "QuantizationConfig | None" = None,
prefix: str = "",
disable_kv_quant: bool = False,
):
super().__init__()
# LTX-2 uses "rms_norm_across_heads", LTX-2.3 uses "rms_norm" -- both
@@ -548,6 +576,8 @@ class LTX2Attention(torch.nn.Module):
head_size=self.head_dim,
total_num_heads=heads,
bias=bias,
quant_config=quant_config,
prefix=f"{prefix}.to_qkv" if prefix else "to_qkv",
)
self.query_num_heads = self.to_qkv.num_heads
self.kv_num_heads = self.to_qkv.num_kv_heads
@@ -562,6 +592,8 @@ class LTX2Attention(torch.nn.Module):
bias=bias,
gather_output=False,
return_bias=False,
quant_config=quant_config,
prefix=f"{prefix}.to_q" if prefix else "to_q",
)
self.to_k = ColumnParallelLinear(
self.cross_attention_dim,
@@ -569,6 +601,8 @@ class LTX2Attention(torch.nn.Module):
bias=bias,
gather_output=False,
return_bias=False,
quant_config=quant_config,
prefix=f"{prefix}.to_k" if prefix else "to_k",
)
self.to_v = ColumnParallelLinear(
self.cross_attention_dim,
@@ -576,6 +610,8 @@ class LTX2Attention(torch.nn.Module):
bias=bias,
gather_output=False,
return_bias=False,
quant_config=quant_config,
prefix=f"{prefix}.to_v" if prefix else "to_v",
)
self.heads = self.query_num_heads
@@ -617,6 +653,8 @@ class LTX2Attention(torch.nn.Module):
bias=out_bias,
input_is_parallel=True,
return_bias=False,
quant_config=quant_config,
prefix=f"{prefix}.to_out.0" if prefix else "to_out.0",
),
torch.nn.Dropout(dropout) if dropout > 0 else torch.nn.Identity(),
]
@@ -627,9 +665,12 @@ class LTX2Attention(torch.nn.Module):
num_kv_heads=self.kv_num_heads,
softmax_scale=1.0 / (dim_head**0.5),
causal=False,
prefix=prefix,
disable_kv_quant=disable_kv_quant,
)
# LTX-2.3: per-head gated attention
# leave unquantized for this linear
if apply_gated_attention:
self.to_gate_logits = nn.Linear(query_dim, heads, bias=True)
else:
@@ -748,6 +789,8 @@ class LTX2VideoTransformerBlock(nn.Module):
elementwise_affine: bool = False,
rope_type: str = "interleaved",
perturbed_attn: bool = False,
quant_config: "QuantizationConfig | None" = None,
prefix: str = "",
):
super().__init__()
self.video_cross_attn_adaln = video_cross_attn_adaln
@@ -767,6 +810,8 @@ class LTX2VideoTransformerBlock(nn.Module):
qk_norm=qk_norm,
rope_type=rope_type,
apply_gated_attention=video_gated_attn,
quant_config=quant_config,
prefix=f"{prefix}.attn1" if prefix else "attn1",
)
self.audio_norm1 = _make_rms_norm(audio_dim, eps=eps, elementwise_affine=elementwise_affine)
@@ -781,6 +826,8 @@ class LTX2VideoTransformerBlock(nn.Module):
qk_norm=qk_norm,
rope_type=rope_type,
apply_gated_attention=audio_gated_attn,
quant_config=quant_config,
prefix=f"{prefix}.audio_attn1" if prefix else "audio_attn1",
)
# 2. Prompt Cross-Attention
@@ -796,6 +843,9 @@ class LTX2VideoTransformerBlock(nn.Module):
qk_norm=qk_norm,
rope_type=rope_type,
apply_gated_attention=video_gated_attn,
quant_config=quant_config,
prefix=f"{prefix}.attn2" if prefix else "attn2",
disable_kv_quant=True,
)
self.audio_norm2 = _make_rms_norm(audio_dim, eps=eps, elementwise_affine=elementwise_affine)
@@ -810,9 +860,14 @@ class LTX2VideoTransformerBlock(nn.Module):
qk_norm=qk_norm,
rope_type=rope_type,
apply_gated_attention=audio_gated_attn,
quant_config=quant_config,
prefix=f"{prefix}.audio_attn2" if prefix else "audio_attn2",
disable_kv_quant=True,
)
# 3. Audio-to-Video (a2v) and Video-to-Audio (v2a) Cross-Attention
# K/V here come from the other modality's latents (long sequences),
# so KV-cache quant remains enabled.
self.audio_to_video_norm = _make_rms_norm(dim, eps=eps, elementwise_affine=elementwise_affine)
self.audio_to_video_attn = LTX2Attention(
query_dim=dim,
@@ -825,6 +880,8 @@ class LTX2VideoTransformerBlock(nn.Module):
qk_norm=qk_norm,
rope_type=rope_type,
apply_gated_attention=video_gated_attn,
quant_config=quant_config,
prefix=f"{prefix}.audio_to_video_attn" if prefix else "audio_to_video_attn",
)
self.video_to_audio_norm = _make_rms_norm(audio_dim, eps=eps, elementwise_affine=elementwise_affine)
@@ -839,14 +896,26 @@ class LTX2VideoTransformerBlock(nn.Module):
qk_norm=qk_norm,
rope_type=rope_type,
apply_gated_attention=audio_gated_attn,
quant_config=quant_config,
prefix=f"{prefix}.video_to_audio_attn" if prefix else "video_to_audio_attn",
)
# 4. Feedforward layers
self.norm3 = _make_rms_norm(dim, eps=eps, elementwise_affine=elementwise_affine)
self.ff = LTX2FeedForward(dim, activation_fn=activation_fn)
self.ff = LTX2FeedForward(
dim,
activation_fn=activation_fn,
quant_config=quant_config,
prefix=f"{prefix}.ff" if prefix else "ff",
)
self.audio_norm3 = _make_rms_norm(audio_dim, eps=eps, elementwise_affine=elementwise_affine)
self.audio_ff = LTX2FeedForward(audio_dim, activation_fn=activation_fn)
self.audio_ff = LTX2FeedForward(
audio_dim,
activation_fn=activation_fn,
quant_config=quant_config,
prefix=f"{prefix}.audio_ff" if prefix else "audio_ff",
)
# 5. Per-Layer Modulation Parameters
# LTX-2.3 with cross_attn_adaln uses 9 params (extra 3 for cross-attn modulation);
@@ -1365,6 +1434,9 @@ class LTX2VideoTransformer3DModel(nn.Module):
_layerwise_offload_blocks_attrs = ["transformer_blocks"]
_hsdp_shard_conditions = [is_transformer_block_module]
_sp_plan: dict[str, Any] | None = None
packed_modules_mapping = {
"to_qkv": ["to_q", "to_k", "to_v"],
}
@staticmethod
def _build_sp_plan(rope_type: str) -> dict[str, Any]:
@@ -1453,6 +1525,7 @@ class LTX2VideoTransformer3DModel(nn.Module):
cross_attn_mod: bool = False,
audio_gated_attn: bool = False,
audio_cross_attn_mod: bool = False,
quant_config: "QuantizationConfig | None" = None,
) -> None:
super().__init__()
self.perturbed_attn = perturbed_attn
@@ -1647,8 +1720,10 @@ class LTX2VideoTransformer3DModel(nn.Module):
elementwise_affine=norm_elementwise_affine,
rope_type=rope_type,
perturbed_attn=perturbed_attn,
quant_config=quant_config,
prefix=f"transformer_blocks.{layer_idx}",
)
for _ in range(num_layers)
for layer_idx in range(num_layers)
]
)
vllm_omni/diffusion/models/ltx2/pipeline_ltx2.py
+14 -4
View File
@@ -9,10 +9,13 @@ import json
import os
from collections.abc import Iterable
from contextlib import nullcontext
from typing import Any, ClassVar
from typing import TYPE_CHECKING, Any, ClassVar
import numpy as np
import torch
if TYPE_CHECKING:
from vllm.model_executor.layers.quantization.base_config import QuantizationConfig
from diffusers import AutoencoderKLLTX2Audio, AutoencoderKLLTX2Video, FlowMatchEulerDiscreteScheduler
from diffusers.pipelines.ltx2 import LTX2TextConnectors
from diffusers.pipelines.ltx2.utils import DISTILLED_SIGMA_VALUES, STAGE_2_DISTILLED_SIGMA_VALUES
@@ -67,14 +70,20 @@ def load_transformer_config(model_path: str, subfolder: str = "transformer", loc
return {}
def create_transformer_from_config(config: dict) -> LTX2VideoTransformer3DModel:
def create_transformer_from_config(
config: dict,
quant_config: QuantizationConfig | None = None,
) -> LTX2VideoTransformer3DModel:
"""Create LTX2VideoTransformer3DModel from config dict."""
if not config:
if not config and quant_config is None:
return LTX2VideoTransformer3DModel()
signature = inspect.signature(LTX2VideoTransformer3DModel.__init__)
allowed_keys = set(signature.parameters.keys())
kwargs = {k: v for k, v in config.items() if k in allowed_keys}
if quant_config is not None:
kwargs["quant_config"] = quant_config
return LTX2VideoTransformer3DModel(**kwargs)
@@ -214,7 +223,8 @@ class LTX2Pipeline(nn.Module, CFGParallelMixin, ProgressBarMixin):
).to(self.device)
transformer_config = load_transformer_config(model, "transformer", local_files_only)
self.transformer = create_transformer_from_config(transformer_config)
quant_config = getattr(self.od_config, "quantization_config", None)
self.transformer = create_transformer_from_config(transformer_config, quant_config=quant_config)
self.scheduler = FlowMatchEulerDiscreteScheduler.from_pretrained(
model, subfolder="scheduler", local_files_only=local_files_only
vllm_omni/diffusion/models/ltx2/pipeline_ltx2_3.py
+2 -1
View File
@@ -174,7 +174,8 @@ class LTX23Pipeline(nn.Module, ProgressBarMixin):
# --- Transformer: created empty, weights loaded via AutoWeightsLoader ---
transformer_config = load_transformer_config(model, "transformer", local_files_only)
self.transformer = create_transformer_from_config(transformer_config)
quant_config = getattr(self.od_config, "quantization_config", None)
self.transformer = create_transformer_from_config(transformer_config, quant_config=quant_config)
# --- Scheduler ---
self.scheduler = FlowMatchEulerDiscreteScheduler.from_pretrained(
vllm_omni/diffusion/models/qwen_image/pipeline_qwen_image.py
+16 -1
View File
@@ -290,7 +290,22 @@ class QwenImagePipeline(nn.Module, QwenImageCFGParallelMixin, DiffusionPipelineP
)
self.text_encoder = Qwen2_5_VLForConditionalGeneration.from_pretrained(
model, subfolder="text_encoder", local_files_only=local_files_only
).to(self.device)
)
# Qwen2.5-VL ships a vision tower that text-to-image does not use.
# Drop it while the model is still on CPU, before moving to GPU, so
# the vision tower never consumes GPU memory. Handle both transformers
# layouts: newer puts visual under .model, older puts it directly on
# the model.
visual_owner = None
if hasattr(self.text_encoder, "model") and hasattr(self.text_encoder.model, "visual"):
visual_owner = self.text_encoder.model
elif hasattr(self.text_encoder, "visual"):
visual_owner = self.text_encoder
if visual_owner is not None:
del visual_owner.visual
else:
logger.warning("Qwen-Image: vision tower not found on text encoder; skipping drop")
self.text_encoder = self.text_encoder.to(self.device)
self.vae = DistributedAutoencoderKLQwenImage.from_pretrained(
model, subfolder="vae", local_files_only=local_files_only
).to(self.device)
vllm_omni/diffusion/models/wan2_2/__init__.py
-8
View File
@@ -22,11 +22,6 @@ from .pipeline_wan2_2_s2v import (
get_wan22_s2v_post_process_func,
get_wan22_s2v_pre_process_func,
)
from .pipeline_wan2_2_ti2v import (
Wan22TI2VPipeline,
get_wan22_ti2v_post_process_func,
get_wan22_ti2v_pre_process_func,
)
from .pipeline_wan2_2_vace import (
Wan22VACEPipeline,
get_wan22_vace_post_process_func,
@@ -50,9 +45,6 @@ __all__ = [
"Wan22S2VPipeline",
"get_wan22_s2v_post_process_func",
"get_wan22_s2v_pre_process_func",
"Wan22TI2VPipeline",
"get_wan22_ti2v_post_process_func",
"get_wan22_ti2v_pre_process_func",
"Wan22VACEPipeline",
"get_wan22_vace_post_process_func",
"get_wan22_vace_pre_process_func",
vllm_omni/diffusion/models/wan2_2/pipeline_wan2_2.py
+19 -53
View File
@@ -17,10 +17,12 @@ from torch import nn
from transformers import AutoTokenizer, UMT5EncoderModel
from vllm.model_executor.layers.quantization.base_config import QuantizationConfig
from vllm.model_executor.models.utils import AutoWeightsLoader
from vllm.sequence import IntermediateTensors
from vllm_omni.diffusion.data import DiffusionOutput, OmniDiffusionConfig
from vllm_omni.diffusion.distributed.autoencoders.autoencoder_kl_wan import DistributedAutoencoderKLWan
from vllm_omni.diffusion.distributed.cfg_parallel import CFGParallelMixin
from vllm_omni.diffusion.distributed.pipeline_parallel import AsyncLatents, PipelineParallelMixin
from vllm_omni.diffusion.distributed.utils import get_local_device
from vllm_omni.diffusion.forward_context import set_forward_context_denoise_step_idx
from vllm_omni.diffusion.model_loader.diffusers_loader import DiffusersPipelineLoader
@@ -120,8 +122,7 @@ def load_transformer_config(model_path: str, subfolder: str = "transformer", loc
def create_transformer_from_config(
config: dict,
quant_config: QuantizationConfig | None = None,
config: dict, quant_config: QuantizationConfig | None = None, prefix: str = ""
) -> WanTransformer3DModel:
"""Create WanTransformer3DModel from config dict."""
kwargs: dict = {}
@@ -157,57 +158,15 @@ def create_transformer_from_config(
if "pos_embed_seq_len" in config:
kwargs["pos_embed_seq_len"] = config["pos_embed_seq_len"]
# Auto-detect quantization from transformer's config.json when not explicitly provided.
# merge_mxfp8_checkpoint.py injects quantization_config into config.json so that
# offline quantized checkpoints are recognized here without a CLI flag.
if "quantization_config" in config:
from vllm_omni.quantization.factory import build_quant_config
from vllm_omni.quantization.factory import resolve_quant_config_from_disk
disk_qc = config["quantization_config"]
if isinstance(disk_qc, dict) and "quant_method" in disk_qc:
qc_method = disk_qc["quant_method"]
qc_kwargs = {k: v for k, v in disk_qc.items() if k != "quant_method"}
if quant_config is None:
# No CLI flag: full auto-detection.
quant_config = build_quant_config(qc_method, **qc_kwargs)
logger.info(
"Auto-detected quantization from transformer config.json: method=%s kwargs=%s",
qc_method,
qc_kwargs,
)
elif quant_config.get_name() != qc_method:
# The caller supplied a quant_config of a different method than
# what the checkpoint was built with. Loading serialized tensors
# (e.g. MXFP8 weight scales) with the wrong linear method would
# produce corrupt output or a shape mismatch crash. Reject early
# so the user gets a clear message instead of a silent failure.
raise ValueError(
f"Checkpoint config.json declares quant_method={qc_method!r} but the "
f"active quantization config is {quant_config.get_name()!r}. "
"Pass a matching --quantization flag or omit it for auto-detection."
)
elif (
qc_kwargs.get("is_checkpoint_mxfp8_serialized", False)
and hasattr(quant_config, "is_checkpoint_mxfp8_serialized")
and not quant_config.is_checkpoint_mxfp8_serialized
):
# Same method: CLI provided online mode but config.json marks this
# as a pre-quantized offline checkpoint. Switch to offline mode so
# users can pass --quantization mxfp8 without knowing the
# online/offline distinction.
quant_config = build_quant_config(qc_method, **qc_kwargs)
logger.info(
"config.json marks checkpoint as serialized; switching from online to offline MXFP8 mode.",
)
elif isinstance(disk_qc, str) and quant_config is None:
quant_config = build_quant_config(disk_qc)
logger.info(
"Auto-detected quantization from transformer config.json: method=%s",
disk_qc,
)
quant_config = resolve_quant_config_from_disk(quant_config, config["quantization_config"])
if quant_config is not None:
kwargs["quant_config"] = quant_config
if prefix:
kwargs["prefix"] = prefix
return WanTransformer3DModel(**kwargs)
@@ -297,7 +256,9 @@ def get_wan22_pre_process_func(
return pre_process_func
class Wan22Pipeline(nn.Module, CFGParallelMixin, ProgressBarMixin, DiffusionPipelineProfilerMixin):
class Wan22Pipeline(
nn.Module, PipelineParallelMixin, CFGParallelMixin, ProgressBarMixin, DiffusionPipelineProfilerMixin
):
def __init__(
self,
*,
@@ -460,7 +421,7 @@ class Wan22Pipeline(nn.Module, CFGParallelMixin, ProgressBarMixin, DiffusionPipe
attention_kwargs: dict[str, Any],
latent_condition: torch.Tensor | None = None,
first_frame_mask: torch.Tensor | None = None,
) -> torch.Tensor:
) -> torch.Tensor | AsyncLatents:
if attention_kwargs is None:
attention_kwargs = {}
with self.progress_bar(total=len(timesteps)) as pbar:
@@ -858,7 +819,11 @@ class Wan22Pipeline(nn.Module, CFGParallelMixin, ProgressBarMixin, DiffusionPipe
output=output, stage_durations=self.stage_durations if hasattr(self, "stage_durations") else None
)
def predict_noise(self, current_model: nn.Module | None = None, **kwargs: Any) -> torch.Tensor:
def predict_noise(
self,
current_model: nn.Module | None = None,
**kwargs: Any,
) -> torch.Tensor | IntermediateTensors:
"""
Forward pass through transformer to predict noise.
@@ -867,11 +832,12 @@ class Wan22Pipeline(nn.Module, CFGParallelMixin, ProgressBarMixin, DiffusionPipe
**kwargs: Arguments to pass to the transformer
Returns:
Predicted noise tensor
Predicted noise tensor or IntermediateTensors on non-last PP stages.
"""
if current_model is None:
current_model = self.transformer
return current_model(**kwargs)[0]
result = current_model(**kwargs)
return result if isinstance(result, IntermediateTensors) else result[0]
def encode_prompt(
self,
vllm_omni/diffusion/models/wan2_2/pipeline_wan2_2_i2v.py
+34 -8
View File
@@ -17,10 +17,12 @@ from diffusers.utils.torch_utils import randn_tensor
from torch import nn
from transformers import AutoTokenizer, CLIPImageProcessor, CLIPVisionModel, UMT5EncoderModel
from vllm.model_executor.models.utils import AutoWeightsLoader
from vllm.sequence import IntermediateTensors
from vllm_omni.diffusion.data import DiffusionOutput, OmniDiffusionConfig
from vllm_omni.diffusion.distributed.autoencoders.autoencoder_kl_wan import DistributedAutoencoderKLWan
from vllm_omni.diffusion.distributed.cfg_parallel import CFGParallelMixin
from vllm_omni.diffusion.distributed.pipeline_parallel import AsyncLatents, PipelineParallelMixin
from vllm_omni.diffusion.distributed.utils import get_local_device
from vllm_omni.diffusion.forward_context import set_forward_context_denoise_step_idx
from vllm_omni.diffusion.model_loader.diffusers_loader import DiffusersPipelineLoader
@@ -135,7 +137,12 @@ def get_wan22_i2v_pre_process_func(
class Wan22I2VPipeline(
nn.Module, SupportImageInput, CFGParallelMixin, ProgressBarMixin, DiffusionPipelineProfilerMixin
nn.Module,
SupportImageInput,
PipelineParallelMixin,
CFGParallelMixin,
ProgressBarMixin,
DiffusionPipelineProfilerMixin,
):
"""
Wan2.2 Image-to-Video Pipeline.
@@ -224,10 +231,25 @@ class Wan22I2VPipeline(
# Transformers (weights loaded via load_weights)
# Load config from model directory or HF Hub to get correct in_channels for I2V models
transformer_config = load_transformer_config(model, "transformer", local_files_only)
self.transformer = self._create_transformer(transformer_config)
self.transformer = create_transformer_from_config(
transformer_config,
quant_config=od_config.quantization_config,
)
if self.has_transformer_2:
transformer_2_config = load_transformer_config(model, "transformer_2", local_files_only)
self.transformer_2 = self._create_transformer(transformer_2_config)
t2_quant = transformer_2_config.get("quantization_config")
if isinstance(t2_quant, dict) and "quant_method" in t2_quant:
from vllm_omni.quantization.factory import build_quant_config
method = t2_quant["quant_method"]
kwargs = {k: v for k, v in t2_quant.items() if k != "quant_method"}
t2_quant = build_quant_config(method, **kwargs)
else:
t2_quant = None
self.transformer_2 = create_transformer_from_config(
transformer_2_config,
quant_config=t2_quant,
)
else:
self.transformer_2 = None
@@ -283,7 +305,7 @@ class Wan22I2VPipeline(
attention_kwargs: dict[str, Any],
condition: torch.Tensor,
first_frame_mask: torch.Tensor,
) -> torch.Tensor:
) -> torch.Tensor | AsyncLatents:
if attention_kwargs is None:
attention_kwargs = {}
with self.progress_bar(total=len(timesteps)) as pbar:
@@ -348,7 +370,6 @@ class Wan22I2VPipeline(
# Compute the previous noisy sample x_t -> x_t-1 with automatic CFG sync
latents = self.scheduler_step_maybe_with_cfg(noise_pred, t, latents, do_true_cfg)
pbar.update()
return latents
@@ -667,7 +688,11 @@ class Wan22I2VPipeline(
output=output, stage_durations=self.stage_durations if hasattr(self, "stage_durations") else None
)
def predict_noise(self, current_model: nn.Module | None = None, **kwargs: Any) -> torch.Tensor:
def predict_noise(
self,
current_model: nn.Module | None = None,
**kwargs: Any,
) -> torch.Tensor | IntermediateTensors:
"""
Forward pass through transformer to predict noise.
@@ -676,11 +701,12 @@ class Wan22I2VPipeline(
**kwargs: Arguments to pass to the transformer
Returns:
Predicted noise tensor
Predicted noise tensor or IntermediateTensors on non-last PP stages.
"""
if current_model is None:
current_model = self.transformer
return current_model(**kwargs)[0]
result = current_model(**kwargs)
return result if isinstance(result, IntermediateTensors) else result[0]
def encode_prompt(
self,
vllm_omni/diffusion/models/wan2_2/pipeline_wan2_2_ti2v.py
-731
View File
@@ -1,731 +0,0 @@
# SPDX-License-Identifier: Apache-2.0
# SPDX-FileCopyrightText: Copyright contributors to the vLLM project
"""
Wan2.2 TI2V (Text-Image-to-Video) Pipeline.
This pipeline supports the unified TI2V-5B model that can generate videos from:
- Text only (T2V mode)
- Text + Image (I2V mode)
The key difference from the MoE-based I2V pipeline is:
- Single transformer (not MoE with two transformers)
- Uses expand_timesteps mode for image conditioning
- No CLIP image encoder - only VAE encoding for image condition
"""
from __future__ import annotations
import logging
import os
from collections.abc import Iterable
from typing import Any, cast
import numpy as np
import PIL.Image
import torch
from diffusers.utils.torch_utils import randn_tensor
from torch import nn
from transformers import AutoTokenizer, UMT5EncoderModel
from vllm.model_executor.models.utils import AutoWeightsLoader
from vllm_omni.diffusion.data import DiffusionOutput, OmniDiffusionConfig
from vllm_omni.diffusion.distributed.autoencoders.autoencoder_kl_wan import OmniAutoencoderKLWan
from vllm_omni.diffusion.distributed.cfg_parallel import CFGParallelMixin
from vllm_omni.diffusion.distributed.utils import get_local_device
from vllm_omni.diffusion.forward_context import set_forward_context_denoise_step_idx
from vllm_omni.diffusion.model_loader.diffusers_loader import DiffusersPipelineLoader
from vllm_omni.diffusion.model_loader.hub_prefetch import prefetch_subfolders
from vllm_omni.diffusion.models.interface import SupportImageInput
from vllm_omni.diffusion.models.progress_bar import ProgressBarMixin
from vllm_omni.diffusion.models.wan2_2.pipeline_wan2_2 import (
build_wan_scheduler,
create_transformer_from_config,
load_transformer_config,
resolve_wan_flow_shift,
resolve_wan_sample_solver,
retrieve_latents,
)
from vllm_omni.diffusion.models.wan2_2.wan2_2_transformer import WanTransformer3DModel
from vllm_omni.diffusion.postprocess import interpolate_video_tensor
from vllm_omni.diffusion.request import OmniDiffusionRequest
from vllm_omni.inputs.data import OmniTextPrompt
from vllm_omni.platforms import current_omni_platform
logger = logging.getLogger(__name__)
def get_wan22_ti2v_post_process_func(
od_config: OmniDiffusionConfig,
):
from diffusers.video_processor import VideoProcessor
video_processor = VideoProcessor(vae_scale_factor=8)
def post_process_func(
video: torch.Tensor,
output_type: str = "np",
sampling_params=None,
):
if output_type == "latent":
return video
custom_output = {}
if sampling_params is not None and getattr(sampling_params, "enable_frame_interpolation", False):
video, multiplier = interpolate_video_tensor(
video,
exp=sampling_params.frame_interpolation_exp,
scale=sampling_params.frame_interpolation_scale,
model_path=sampling_params.frame_interpolation_model_path,
)
custom_output["video_fps_multiplier"] = multiplier
return {
"video": video_processor.postprocess_video(video, output_type=output_type),
"custom_output": custom_output,
}
return post_process_func
def get_wan22_ti2v_pre_process_func(
od_config: OmniDiffusionConfig,
):
"""Pre-process function for TI2V: optionally load and resize input image."""
def pre_process_func(request: OmniDiffusionRequest) -> OmniDiffusionRequest:
for i, prompt in enumerate(request.prompts):
multi_modal_data = prompt.get("multi_modal_data", {}) if not isinstance(prompt, str) else None
raw_image = multi_modal_data.get("image", None) if multi_modal_data is not None else None
if isinstance(prompt, str):
prompt = OmniTextPrompt(prompt=prompt)
if "additional_information" not in prompt:
prompt["additional_information"] = {}
if raw_image is None:
raise ValueError(
"""No image is provided. This model requires an image to run.""",
"""Please correctly set `"multi_modal_data": {"image": <an image object or file path>, …}`""",
)
if not isinstance(raw_image, (str, PIL.Image.Image)):
raise TypeError(
f"""Unsupported image format {raw_image.__class__}.""",
"""Please correctly set `"multi_modal_data": {"image": <an image object or file path>, …}`""",
)
image = PIL.Image.open(raw_image).convert("RGB") if isinstance(raw_image, str) else raw_image
# Calculate dimensions based on aspect ratio if not provided
if request.sampling_params.height is None or request.sampling_params.width is None:
# Default max area for 720P (TI2V-5B default)
max_area = 720 * 1280
aspect_ratio = image.height / image.width
# Calculate dimensions maintaining aspect ratio
mod_value = 16 # Must be divisible by 16
height = round(np.sqrt(max_area * aspect_ratio)) // mod_value * mod_value
width = round(np.sqrt(max_area / aspect_ratio)) // mod_value * mod_value
if request.sampling_params.height is None:
request.sampling_params.height = height
if request.sampling_params.width is None:
request.sampling_params.width = width
# Resize image to target dimensions
image = image.resize(
(request.sampling_params.width, request.sampling_params.height), # type: ignore # Above has ensured that width & height are not None
PIL.Image.Resampling.LANCZOS,
)
prompt["multi_modal_data"]["image"] = image # type: ignore # key existence already checked above
request.prompts[i] = prompt
return request
return pre_process_func
class Wan22TI2VPipeline(nn.Module, SupportImageInput, CFGParallelMixin, ProgressBarMixin):
"""
Wan2.2 Text-Image-to-Video (TI2V) Pipeline.
This is a unified pipeline that supports both:
- Text-to-Video (T2V): when no image is provided
- Image-to-Video (I2V): when an image is provided
Uses expand_timesteps mode for I2V conditioning where the first frame
is conditioned on the input image latent.
"""
def __init__(
self,
*,
od_config: OmniDiffusionConfig,
prefix: str = "",
):
super().__init__()
self.od_config = od_config
self.device = get_local_device()
dtype = getattr(od_config, "dtype", torch.bfloat16)
model = od_config.model
local_files_only = os.path.exists(model)
# Set up weights sources for single transformer
self.weights_sources = [
DiffusersPipelineLoader.ComponentSource(
model_or_path=od_config.model,
subfolder="transformer",
revision=None,
prefix="transformer.",
fall_back_to_pt=True,
),
]
# See ``hub_prefetch.py`` for the transformers v5 subfolder race.
prefetch_subfolders(
model,
["tokenizer", "text_encoder", "vae"],
local_files_only=local_files_only,
)
# Text encoder
self.tokenizer = AutoTokenizer.from_pretrained(model, subfolder="tokenizer", local_files_only=local_files_only)
self.text_encoder = UMT5EncoderModel.from_pretrained(
model, subfolder="text_encoder", torch_dtype=dtype, local_files_only=local_files_only
).to(self.device)
# VAE
self.vae = OmniAutoencoderKLWan.from_pretrained(
model, subfolder="vae", torch_dtype=dtype, local_files_only=local_files_only
).to(self.device)
# Single transformer (TI2V uses dense 5B model, not MoE)
# Load config from model to get correct dimensions
transformer_config = load_transformer_config(model, "transformer", local_files_only)
self.transformer = self._create_transformer(transformer_config)
self._sample_solver = "unipc"
self._flow_shift = od_config.flow_shift if od_config.flow_shift is not None else 5.0
self.scheduler = build_wan_scheduler(self._sample_solver, self._flow_shift)
# VAE scale factors
self.vae_scale_factor_temporal = self.vae.config.scale_factor_temporal if hasattr(self.vae, "config") else 4
self.vae_scale_factor_spatial = self.vae.config.scale_factor_spatial if hasattr(self.vae, "config") else 8
# TI2V always uses expand_timesteps mode
self.expand_timesteps = True
self._guidance_scale = None
self._num_timesteps = None
self._current_timestep = None
def _create_transformer(self, config: dict) -> WanTransformer3DModel:
"""Create a transformer from a config dict. Respects od_config.quantization_config."""
quant_config = getattr(self.od_config, "quantization_config", None)
return create_transformer_from_config(config, quant_config=quant_config)
@property
def guidance_scale(self):
return self._guidance_scale
@property
def do_classifier_free_guidance(self):
return self._guidance_scale is not None and self._guidance_scale > 1.0
@property
def num_timesteps(self):
return self._num_timesteps
@property
def current_timestep(self):
return self._current_timestep
def diffuse(
self,
latents: torch.Tensor,
timesteps: torch.Tensor,
prompt_embeds: torch.Tensor,
negative_prompt_embeds: torch.Tensor | None,
guidance_scale: float,
dtype: torch.dtype,
attention_kwargs: dict[str, Any],
num_latent_frames: int,
latent_height: int,
latent_width: int,
latent_condition: torch.Tensor | None = None,
first_frame_mask: torch.Tensor | None = None,
) -> torch.Tensor:
if attention_kwargs is None:
attention_kwargs = {}
with self.progress_bar(total=len(timesteps)) as pbar:
for step_idx, t in enumerate(timesteps):
self._current_timestep = t
set_forward_context_denoise_step_idx(step_idx)
# Prepare latent input
if latent_condition is not None:
# I2V mode: blend condition with latents using mask
latent_model_input = (1 - first_frame_mask) * latent_condition + first_frame_mask * latents
latent_model_input = latent_model_input.to(dtype)
# Expand timesteps for each patch (TI2V style)
temp_ts = (first_frame_mask[0][0][:, ::2, ::2] * t).flatten()
timestep = temp_ts.unsqueeze(0).expand(latents.shape[0], -1)
else:
# T2V mode: use latents directly
latent_model_input = latents.to(dtype)
# Expand timesteps for TI2V model architecture
mask = torch.ones(1, 1, num_latent_frames, latent_height, latent_width, device=latents.device)
temp_ts = (mask[0][0][:, ::2, ::2] * t).flatten()
timestep = temp_ts.unsqueeze(0).expand(latents.shape[0], -1)
do_true_cfg = guidance_scale > 1.0 and negative_prompt_embeds is not None
positive_kwargs = {
"hidden_states": latent_model_input,
"timestep": timestep,
"encoder_hidden_states": prompt_embeds,
"attention_kwargs": attention_kwargs,
"return_dict": False,
"current_model": self.transformer,
}
if do_true_cfg:
negative_kwargs = {
"hidden_states": latent_model_input,
"timestep": timestep,
"encoder_hidden_states": negative_prompt_embeds,
"attention_kwargs": attention_kwargs,
"return_dict": False,
"current_model": self.transformer,
}
else:
negative_kwargs = None
noise_pred = self.predict_noise_maybe_with_cfg(
do_true_cfg=do_true_cfg,
true_cfg_scale=guidance_scale,
positive_kwargs=positive_kwargs,
negative_kwargs=negative_kwargs,
cfg_normalize=False,
)
latents = self.scheduler_step_maybe_with_cfg(noise_pred, t, latents, do_true_cfg)
pbar.update()
return latents
def forward(
self,
req: OmniDiffusionRequest,
prompt: str | None = None,
negative_prompt: str | None = None,
image: PIL.Image.Image | torch.Tensor | None = None,
height: int = 704,
width: int = 1280,
num_inference_steps: int = 40,
guidance_scale: float = 5.0,
frame_num: int = 81,
output_type: str | None = "np",
generator: torch.Generator | list[torch.Generator] | None = None,
prompt_embeds: torch.Tensor | None = None,
negative_prompt_embeds: torch.Tensor | None = None,
attention_kwargs: dict | None = None,
**kwargs,
) -> DiffusionOutput:
# Get parameters from request or arguments
if len(req.prompts) > 1:
raise ValueError(
"""This model only supports a single prompt, not a batched request.""",
"""Please pass in a single prompt object or string, or a single-item list.""",
)
if len(req.prompts) == 1: # If req.prompt is empty, default to prompt & neg_prompt in param list
prompt = req.prompts[0] if isinstance(req.prompts[0], str) else req.prompts[0].get("prompt")
negative_prompt = None if isinstance(req.prompts[0], str) else req.prompts[0].get("negative_prompt")
if prompt is None and prompt_embeds is None:
raise ValueError("Prompt or prompt_embeds is required for Wan2.2 generation.")
# Get image from request (optional for TI2V)
if image is None:
multi_modal_data = (
req.prompts[0].get("multi_modal_data", {}) if not isinstance(req.prompts[0], str) else None
)
raw_image = multi_modal_data.get("image", None) if multi_modal_data is not None else None
if isinstance(raw_image, list):
if len(raw_image) > 1:
logger.warning(
"""Received a list of image. Only a single image is supported by this model."""
"""Taking only the first image for now."""
)
raw_image = raw_image[0]
if raw_image is None:
image = None
elif isinstance(raw_image, str):
image = PIL.Image.open(raw_image)
else:
image = cast(PIL.Image.Image | torch.Tensor, raw_image)
# Default dimensions for TI2V-5B (720P)
height = req.sampling_params.height or height
width = req.sampling_params.width or width
num_frames = req.sampling_params.num_frames if req.sampling_params.num_frames else frame_num
num_steps = req.sampling_params.num_inference_steps or num_inference_steps
# Respect per-request guidance_scale when explicitly provided.
if req.sampling_params.guidance_scale_provided:
guidance_scale = req.sampling_params.guidance_scale
self._guidance_scale = guidance_scale
# Validate inputs
self.check_inputs(
prompt=prompt,
negative_prompt=negative_prompt,
image=image,
height=height,
width=width,
prompt_embeds=prompt_embeds,
negative_prompt_embeds=negative_prompt_embeds,
)
# Adjust num_frames to be compatible with VAE temporal scaling
if num_frames % self.vae_scale_factor_temporal != 1:
num_frames = num_frames // self.vae_scale_factor_temporal * self.vae_scale_factor_temporal + 1
num_frames = max(num_frames, 1)
device = self.device
dtype = self.transformer.dtype
# Generator setup
if generator is None:
generator = req.sampling_params.generator
if generator is None and req.sampling_params.seed is not None:
generator = torch.Generator(device=device).manual_seed(req.sampling_params.seed)
# Encode prompts
if prompt_embeds is None:
prompt_embeds, negative_prompt_embeds = self.encode_prompt(
prompt=prompt,
negative_prompt=negative_prompt,
do_classifier_free_guidance=guidance_scale > 1.0,
num_videos_per_prompt=req.sampling_params.num_outputs_per_prompt or 1,
max_sequence_length=req.sampling_params.max_sequence_length or 512,
device=device,
dtype=dtype,
)
else:
prompt_embeds = prompt_embeds.to(device=device, dtype=dtype)
if negative_prompt_embeds is not None:
negative_prompt_embeds = negative_prompt_embeds.to(device=device, dtype=dtype)
batch_size = prompt_embeds.shape[0]
sample_solver = resolve_wan_sample_solver(req, default=self._sample_solver)
flow_shift = resolve_wan_flow_shift(req, self.od_config)
if sample_solver != self._sample_solver or abs(flow_shift - self._flow_shift) > 1e-6:
self.scheduler = build_wan_scheduler(sample_solver, flow_shift)
self._sample_solver = sample_solver
self._flow_shift = flow_shift
# Timesteps
self.scheduler.set_timesteps(num_steps, device=device)
timesteps = self.scheduler.timesteps
self._num_timesteps = len(timesteps)
# Prepare latents
num_channels_latents = self.transformer.config.in_channels
num_latent_frames = (num_frames - 1) // self.vae_scale_factor_temporal + 1
latent_height = height // self.vae_scale_factor_spatial
latent_width = width // self.vae_scale_factor_spatial
# Check if we have an image (I2V mode) or not (T2V mode)
if image is not None:
# I2V mode: prepare latents with image condition
from diffusers.video_processor import VideoProcessor
video_processor = VideoProcessor(vae_scale_factor=self.vae_scale_factor_spatial)
if isinstance(image, PIL.Image.Image):
image_tensor = video_processor.preprocess(image, height=height, width=width)
else:
image_tensor = image
image_tensor = image_tensor.to(device=device, dtype=torch.float32)
latents, latent_condition, first_frame_mask = self.prepare_i2v_latents(
image=image_tensor,
batch_size=batch_size,
num_channels_latents=num_channels_latents,
height=height,
width=width,
num_frames=num_frames,
dtype=torch.float32,
device=device,
generator=generator,
latents=req.sampling_params.latents,
)
else:
# T2V mode: prepare random latents
latents = self.prepare_latents(
batch_size=batch_size,
num_channels_latents=num_channels_latents,
height=height,
width=width,
num_frames=num_frames,
dtype=torch.float32,
device=device,
generator=generator,
latents=req.sampling_params.latents,
)
latent_condition = None
first_frame_mask = torch.ones(
1, 1, num_latent_frames, latent_height, latent_width, dtype=torch.float32, device=device
)
if attention_kwargs is None:
attention_kwargs = {}
latents = self.diffuse(
latents=latents,
timesteps=timesteps,
prompt_embeds=prompt_embeds,
negative_prompt_embeds=negative_prompt_embeds,
guidance_scale=guidance_scale,
dtype=dtype,
attention_kwargs=attention_kwargs,
num_latent_frames=num_latent_frames,
latent_height=latent_height,
latent_width=latent_width,
latent_condition=latent_condition,
first_frame_mask=first_frame_mask,
)
# Wan2.2 is prone to out of memory errors when predicting large videos
# so we empty the cache here to avoid OOM before vae decoding.
if current_omni_platform.is_available():
current_omni_platform.empty_cache()
self._current_timestep = None
# For I2V mode, blend final latents with condition
if latent_condition is not None:
latents = (1 - first_frame_mask) * latent_condition + first_frame_mask * latents
# Decode
if output_type == "latent":
output = latents
else:
latents = latents.to(self.vae.dtype)
latents_mean = (
torch.tensor(self.vae.config.latents_mean)
.view(1, self.vae.config.z_dim, 1, 1, 1)
.to(latents.device, latents.dtype)
)
latents_std = 1.0 / torch.tensor(self.vae.config.latents_std).view(1, self.vae.config.z_dim, 1, 1, 1).to(
latents.device, latents.dtype
)
latents = latents / latents_std + latents_mean
output = self.vae.decode(latents, return_dict=False)[0]
return DiffusionOutput(output=output)
def predict_noise(self, current_model: nn.Module | None = None, **kwargs: Any) -> torch.Tensor:
"""
Forward pass through transformer to predict noise.
Args:
current_model: The transformer model to use
**kwargs: Arguments to pass to the transformer
Returns:
Predicted noise tensor
"""
if current_model is None:
current_model = self.transformer
return current_model(**kwargs)[0]
def encode_prompt(
self,
prompt: str | list[str],
negative_prompt: str | list[str] | None = None,
do_classifier_free_guidance: bool = True,
num_videos_per_prompt: int = 1,
max_sequence_length: int = 512,
device: torch.device | None = None,
dtype: torch.dtype | None = None,
):
"""Encode text prompts using T5 text encoder."""
device = device or self.device
dtype = dtype or self.text_encoder.dtype
prompt = [prompt] if isinstance(prompt, str) else prompt
prompt_clean = [self._prompt_clean(p) for p in prompt]
batch_size = len(prompt_clean)
text_inputs = self.tokenizer(
prompt_clean,
padding="max_length",
max_length=max_sequence_length,
truncation=True,
add_special_tokens=True,
return_attention_mask=True,
return_tensors="pt",
)
ids, mask = text_inputs.input_ids, text_inputs.attention_mask
seq_lens = mask.gt(0).sum(dim=1).long()
prompt_embeds = self.text_encoder(ids.to(device), mask.to(device)).last_hidden_state
prompt_embeds = prompt_embeds.to(dtype=dtype, device=device)
prompt_embeds = [u[:v] for u, v in zip(prompt_embeds, seq_lens)]
prompt_embeds = torch.stack(
[torch.cat([u, u.new_zeros(max_sequence_length - u.size(0), u.size(1))]) for u in prompt_embeds], dim=0
)
_, seq_len, _ = prompt_embeds.shape
prompt_embeds = prompt_embeds.repeat(1, num_videos_per_prompt, 1)
prompt_embeds = prompt_embeds.view(batch_size * num_videos_per_prompt, seq_len, -1)
negative_prompt_embeds = None
if do_classifier_free_guidance:
negative_prompt = negative_prompt or ""
negative_prompt = batch_size * [negative_prompt] if isinstance(negative_prompt, str) else negative_prompt
neg_text_inputs = self.tokenizer(
[self._prompt_clean(p) for p in negative_prompt],
padding="max_length",
max_length=max_sequence_length,
truncation=True,
add_special_tokens=True,
return_attention_mask=True,
return_tensors="pt",
)
ids_neg, mask_neg = neg_text_inputs.input_ids, neg_text_inputs.attention_mask
seq_lens_neg = mask_neg.gt(0).sum(dim=1).long()
negative_prompt_embeds = self.text_encoder(ids_neg.to(device), mask_neg.to(device)).last_hidden_state
negative_prompt_embeds = negative_prompt_embeds.to(dtype=dtype, device=device)
negative_prompt_embeds = [u[:v] for u, v in zip(negative_prompt_embeds, seq_lens_neg)]
negative_prompt_embeds = torch.stack(
[
torch.cat([u, u.new_zeros(max_sequence_length - u.size(0), u.size(1))])
for u in negative_prompt_embeds
],
dim=0,
)
negative_prompt_embeds = negative_prompt_embeds.repeat(1, num_videos_per_prompt, 1)
negative_prompt_embeds = negative_prompt_embeds.view(batch_size * num_videos_per_prompt, seq_len, -1)
return prompt_embeds, negative_prompt_embeds
@staticmethod
def _prompt_clean(text: str) -> str:
return " ".join(text.strip().split())
def prepare_latents(
self,
batch_size: int,
num_channels_latents: int,
height: int,
width: int,
num_frames: int,
dtype: torch.dtype | None,
device: torch.device | None,
generator: torch.Generator | list[torch.Generator] | None,
latents: torch.Tensor | None = None,
) -> torch.Tensor:
"""Prepare random latents for T2V mode."""
if latents is not None:
return latents.to(device=device, dtype=dtype)
num_latent_frames = (num_frames - 1) // self.vae_scale_factor_temporal + 1
shape = (
batch_size,
num_channels_latents,
num_latent_frames,
int(height) // self.vae_scale_factor_spatial,
int(width) // self.vae_scale_factor_spatial,
)
if isinstance(generator, list) and len(generator) != batch_size:
raise ValueError(f"Generator list length {len(generator)} does not match batch size {batch_size}.")
latents = randn_tensor(shape, generator=generator, device=device, dtype=dtype)
return latents
def prepare_i2v_latents(
self,
image: torch.Tensor,
batch_size: int,
num_channels_latents: int,
height: int,
width: int,
num_frames: int,
dtype: torch.dtype | None,
device: torch.device | None,
generator: torch.Generator | list[torch.Generator] | None,
latents: torch.Tensor | None = None,
) -> tuple[torch.Tensor, torch.Tensor, torch.Tensor]:
"""
Prepare latents for I2V mode with image conditioning.
Returns:
latents: Initial noise latents
latent_condition: Encoded first frame condition
first_frame_mask: Mask (0 for first frame, 1 for rest)
"""
num_latent_frames = (num_frames - 1) // self.vae_scale_factor_temporal + 1
latent_height = height // self.vae_scale_factor_spatial
latent_width = width // self.vae_scale_factor_spatial
shape = (batch_size, num_channels_latents, num_latent_frames, latent_height, latent_width)
if latents is None:
latents = randn_tensor(shape, generator=generator, device=device, dtype=dtype)
else:
latents = latents.to(device=device, dtype=dtype)
# Prepare first frame condition
image = image.unsqueeze(2) # [batch, channels, 1, height, width]
image = image.to(device=device, dtype=self.vae.dtype)
# Encode through VAE
latent_condition = retrieve_latents(self.vae.encode(image), sample_mode="argmax")
latent_condition = latent_condition.repeat(batch_size, 1, 1, 1, 1)
# Normalize latents
latents_mean = (
torch.tensor(self.vae.config.latents_mean)
.view(1, self.vae.config.z_dim, 1, 1, 1)
.to(latent_condition.device, latent_condition.dtype)
)
latents_std = 1.0 / torch.tensor(self.vae.config.latents_std).view(1, self.vae.config.z_dim, 1, 1, 1).to(
latent_condition.device, latent_condition.dtype
)
latent_condition = (latent_condition - latents_mean) * latents_std
latent_condition = latent_condition.to(dtype)
# Create mask: 0 for first frame (condition), 1 for rest (to denoise)
first_frame_mask = torch.ones(1, 1, num_latent_frames, latent_height, latent_width, dtype=dtype, device=device)
first_frame_mask[:, :, 0] = 0
return latents, latent_condition, first_frame_mask
def check_inputs(
self,
prompt,
negative_prompt,
image,
height,
width,
prompt_embeds=None,
negative_prompt_embeds=None,
):
if height % 16 != 0 or width % 16 != 0:
raise ValueError(f"`height` and `width` have to be divisible by 16 but are {height} and {width}.")
if prompt is not None and prompt_embeds is not None:
raise ValueError("Cannot forward both `prompt` and `prompt_embeds`. Please provide only one.")
if negative_prompt is not None and negative_prompt_embeds is not None:
raise ValueError(
"Cannot forward both `negative_prompt` and `negative_prompt_embeds`. Please provide only one."
)
if prompt is None and prompt_embeds is None:
raise ValueError("Provide either `prompt` or `prompt_embeds`.")
def load_weights(self, weights: Iterable[tuple[str, torch.Tensor]]) -> set[str]:
"""Load weights using AutoWeightsLoader for vLLM integration."""
loader = AutoWeightsLoader(self)
return loader.load_weights(weights)
vllm_omni/diffusion/models/wan2_2/pipeline_wan2_2_vace.py
+3
View File
@@ -45,6 +45,7 @@ logger = init_logger(__name__)
def create_vace_transformer_from_config(
config: dict,
quant_config: QuantizationConfig | None = None,
prefix: str = "",
) -> WanVACETransformer3DModel:
"""Create WanVACETransformer3DModel from config dict."""
kwargs = {}
@@ -84,6 +85,8 @@ def create_vace_transformer_from_config(
kwargs["vace_in_channels"] = config["vace_in_channels"]
if quant_config is not None:
kwargs["quant_config"] = quant_config
if prefix:
kwargs["prefix"] = prefix
return WanVACETransformer3DModel(**kwargs)
vllm_omni/diffusion/models/wan2_2/wan2_2_transformer.py
+82 -40
View File
@@ -21,9 +21,21 @@ from vllm.model_executor.layers.conv import Conv3dLayer
from vllm.model_executor.layers.linear import ColumnParallelLinear, QKVParallelLinear, RowParallelLinear
from vllm.model_executor.layers.quantization.base_config import QuantizationConfig
from vllm.model_executor.model_loader.weight_utils import default_weight_loader
from vllm.model_executor.models.utils import (
PPMissingLayer,
is_pp_missing_parameter,
make_empty_intermediate_tensors_factory,
make_layers,
)
from vllm.sequence import IntermediateTensors
from vllm_omni.diffusion.attention.backends.abstract import AttentionMetadata
from vllm_omni.diffusion.attention.layer import Attention
from vllm_omni.diffusion.distributed.parallel_state import (
get_pipeline_parallel_world_size,
is_pipeline_first_stage,
is_pipeline_last_stage,
)
from vllm_omni.diffusion.distributed.sp_plan import (
SequenceParallelInput,
SequenceParallelOutput,
@@ -871,14 +883,23 @@ class WanTransformer3DModel(nn.Module):
inner_dim = num_attention_heads * attention_head_dim
out_channels = out_channels or in_channels
if get_pipeline_parallel_world_size() == 0 and not is_pipeline_first_stage():
raise RuntimeError(
"`initialize_model_parallel()` must be called before constructing `WanTransformer3DModel`"
)
# 1. Patch & position embedding
self.rope = WanRotaryPosEmbed(attention_head_dim, patch_size, rope_max_seq_len)
self.patch_embedding = Conv3dLayer(
in_channels=in_channels,
out_channels=inner_dim,
kernel_size=patch_size,
stride=patch_size,
)
# Patch embedding only on first PP stage; other stages receive hidden_states via P2P
if is_pipeline_first_stage():
self.patch_embedding = Conv3dLayer(
in_channels=in_channels,
out_channels=inner_dim,
kernel_size=patch_size,
stride=patch_size,
)
else:
self.patch_embedding = PPMissingLayer()
# 2. Condition embeddings
self.condition_embedder = WanTimeTextImageEmbedding(
@@ -890,30 +911,39 @@ class WanTransformer3DModel(nn.Module):
pos_embed_seq_len=pos_embed_seq_len,
)
# 3. Transformer blocks
self.blocks = nn.ModuleList(
[
WanTransformerBlock(
inner_dim,
ffn_dim,
num_attention_heads,
eps,
added_kv_proj_dim,
cross_attn_norm,
quant_config=quant_config,
prefix=f"blocks.{layer_idx}",
)
for layer_idx in range(num_layers)
]
# 3. Transformer blocks — partitioned across PP stages via vLLM's `make_layers`.
# It computes the [start_layer, end_layer) slice for this rank and fills the remaining slots
# with PPMissingLayer so that weight names stay globally consistent.
self.start_layer, self.end_layer, self.blocks = make_layers(
num_layers,
lambda prefix: WanTransformerBlock(
inner_dim,
ffn_dim,
num_attention_heads,
eps,
added_kv_proj_dim,
cross_attn_norm,
quant_config=quant_config,
prefix=prefix,
),
prefix="blocks",
)
# 4. Output norm & projection
self.norm_out = AdaLayerNorm(inner_dim, elementwise_affine=False, eps=eps)
self.proj_out = nn.Linear(inner_dim, out_channels * math.prod(patch_size))
# 4. Output norm & projection — only on the last PP stage
if is_pipeline_last_stage():
self.norm_out = AdaLayerNorm(inner_dim, elementwise_affine=False, eps=eps)
self.proj_out = nn.Linear(inner_dim, out_channels * math.prod(patch_size))
else:
self.norm_out, self.proj_out = PPMissingLayer(), PPMissingLayer()
# SP helper modules
self.timestep_proj_prepare = TimestepProjPrepare()
self.output_scale_shift_prepare = OutputScaleShiftPrepare(inner_dim)
if is_pipeline_last_stage():
self.output_scale_shift_prepare = OutputScaleShiftPrepare(inner_dim)
else:
self.output_scale_shift_prepare = PPMissingLayer()
self.make_empty_intermediate_tensors = make_empty_intermediate_tensors_factory(["hidden_states"], inner_dim)
# ROPE helper
self._cached_rope_emb = None
@@ -930,9 +960,10 @@ class WanTransformer3DModel(nn.Module):
timestep: torch.LongTensor,
encoder_hidden_states: torch.Tensor,
encoder_hidden_states_image: torch.Tensor | None = None,
intermediate_tensors: IntermediateTensors | None = None,
return_dict: bool = True,
attention_kwargs: dict[str, Any] | None = None,
) -> torch.Tensor | Transformer2DModelOutput:
) -> torch.Tensor | Transformer2DModelOutput | IntermediateTensors:
batch_size, num_channels, num_frames, height, width = hidden_states.shape
p_t, p_h, p_w = self.config.patch_size
post_patch_num_frames = num_frames // p_t
@@ -949,18 +980,25 @@ class WanTransformer3DModel(nn.Module):
self._hidden_states_shape = hidden_states.shape
self._cached_rope_emb = rotary_emb
# Patch embedding and flatten to sequence
# (hidden_states is sharded at blocks.0 input by _sp_plan)
hidden_states = self.patch_embedding(hidden_states)
hidden_states = hidden_states.flatten(2).transpose(1, 2)
if is_pipeline_first_stage():
# Patch embedding and flatten to sequence
# (hidden_states is sharded at blocks.0 input by _sp_plan)
hidden_states = self.patch_embedding(hidden_states)
hidden_states = hidden_states.flatten(2).transpose(1, 2)
else:
if intermediate_tensors is None:
raise RuntimeError("intermediate_tensors must be provided for non-first PP stages")
hidden_states = intermediate_tensors["hidden_states"]
# Handle timestep shape
# Handle timestep shape (2-D for TI2V, 1-D for T2V)
if timestep.ndim == 2:
ts_seq_len = timestep.shape[1]
timestep = timestep.flatten()
else:
ts_seq_len = None
# Compute conditioning on all PP stages.
# Each stage needs temb/timestep_proj for scale-shift modulation in its local blocks.
temb, timestep_proj, encoder_hidden_states, encoder_hidden_states_image = self.condition_embedder(
timestep, encoder_hidden_states, encoder_hidden_states_image, timestep_seq_len=ts_seq_len
)
@@ -996,14 +1034,13 @@ class WanTransformer3DModel(nn.Module):
hidden_states_mask = None
# Transformer blocks
for block in self.blocks:
hidden_states = block(
hidden_states,
encoder_hidden_states,
timestep_proj,
rotary_emb,
hidden_states_mask,
)
for block in self.blocks[self.start_layer : self.end_layer]:
hidden_states = block(hidden_states, encoder_hidden_states, timestep_proj, rotary_emb, hidden_states_mask)
if not is_pipeline_last_stage():
# Non-last PP stage: hand the token sequence to the caller via IntermediateTensors.
# predict_noise will broadcast it to the next stage before calling that stage's forward.
return IntermediateTensors({"hidden_states": hidden_states})
# Output norm, projection & unpatchify
shift, scale = self.output_scale_shift_prepare(temb)
@@ -1074,7 +1111,8 @@ class WanTransformer3DModel(nn.Module):
if weight_name not in original_name:
continue
lookup_name = original_name.replace(weight_name, param_name)
if lookup_name not in params_dict:
# Skip weights that belong to PP stages other than this one
if is_pp_missing_parameter(lookup_name, self) or lookup_name not in params_dict:
break
param = params_dict[lookup_name]
weight_loader = param.weight_loader
@@ -1098,6 +1136,10 @@ class WanTransformer3DModel(nn.Module):
if modulation_alias in params_dict:
lookup_name = modulation_alias
# Skip weights that belong to PP stages other than this one
if is_pp_missing_parameter(lookup_name, self):
continue
if lookup_name not in params_dict:
logger.warning(f"Skipping weight {original_name} -> {lookup_name}")
continue