[New Model]: Support Loading Original (Non-Diffusers) PAI/Wan2.2-VACE-Fun-A14B Checkpoint

The model to consider.

https://modelscope.ai/models/PAI/Wan2.2-VACE-Fun-A14B

The closest model vllm-omni already supports.

No response

What's your difficulty of supporting the model you want?

Summary vllm-omni's Wan22VACEPipeline only supports the diffusers-format Wan2.1-VACE-14B checkpoint (e.g., Wan-AI/Wan2.1-VACE-14B-diffusers on HuggingFace). The original Wan2.1-VACE-14B release uses a different directory layout and weight key naming convention, making it impossible to load directly without manual conversion or code changes. Original Model Directory Structure The original (non-diffusers) Wan2.1-VACE-14B has the following layout: wan2.1-vace-14B/ ├── configuration.json # Meta info, NOT model_index.json ├── Wan2.1_VAE.pth # VAE weights (.pth / pickle format) ├── models_t5_umt5-xxl-enc-bf16.pth # T5 text encoder (.pth / pickle format) ├── google/umt5-xxl/ # T5 tokenizer files │ ├── spiece.model │ ├── tokenizer.json │ ├── tokenizer_config.json │ └── special_tokens_map.json ├── high_noise_model/ # High-noise stage DiT │ ├── config.json # _class_name: VaceWanModel │ └── diffusion_pytorch_model.safetensors └── low_noise_model/ # Low-noise stage DiT ├── config.json # _class_name: VaceWanModel └── diffusion_pytorch_model.safetensors What vllm-omni Expects Wan22Pipeline.init (pipeline_wan2_2.py) expects the diffusers-format layout: model/ ├── model_index.json # With "_class_name": "WanVACEPipeline" ├── tokenizer/ # Loaded via AutoTokenizer.from_pretrained(subfolder="tokenizer") ├── text_encoder/ # Loaded via UMT5EncoderModel.from_pretrained(subfolder="text_encoder") ├── vae/ # Loaded via DistributedAutoencoderKLWan.from_pretrained(subfolder="vae") ├── transformer/ # High-noise DiT, loaded via load_transformer_config(model, "transformer") │ ├── config.json │ └── diffusion_pytorch_model.safetensors └── transformer_2/ # Low-noise DiT, loaded via load_transformer_config(model, "transformer_2") ├── config.json └── diffusion_pytorch_model.safetensors Issue 1: Directory Layout Mismatch The original model lacks model_index.json and uses different subfolder names (high_noise_model/low_noise_model instead of transformer/transformer_2, google/umt5-xxl instead of tokenizer, flat .pth files instead of text_encoder/ and vae/ subdirectories). The config auto-detection in OmniDiffusionConfig.enrich_config() (data.py:696-770) relies on model_index.json to identify the pipeline class. Without it, the code falls back to reading config.json and mapping model_type, but "vace" is not handled — only "bagel", "nextstep", and "s2v" are mapped. Issue 2: Weight Key Name Mismatches (Transformer) The original model's safetensors uses a different parameter naming convention than vllm-omni's WanTransformer3DModel / WanVACETransformer3DModel. The existing load_weights() method (wan2_2_transformer.py:945-1050) handles some remapping (modulation → scale_shift_table, ffn.net.0 → ffn.net_0, QKV fusion for attn1.to_q/k/v), but the following mappings are missing: Original Weight Key vllm-omni Expected Key blocks.N.self_attn.q/k/v/o blocks.N.attn1.to_q/to_k/to_v/to_out blocks.N.self_attn.norm_q/norm_k blocks.N.attn1.norm_q/norm_k blocks.N.cross_attn.q/k/v/o blocks.N.attn2.to_q/to_k/to_v/to_out blocks.N.cross_attn.norm_q/norm_k blocks.N.attn2.norm_q/norm_k text_embedding.* condition_embedder.text_embedder.* time_embedding.* condition_embedder.time_embedder.* time_projection.* condition_embedder.time_proj.* head.head.* proj_out.* head.modulation output_scale_shift_prepare.scale_shift_table Issue 3: VACE Block Weight Name Mismatches The original VACE model uses before_proj/after_proj in vace_blocks, but vllm-omni's VaceWanTransformerBlock (wan2_2_vace_transformer.py:25-63) defines them as proj_in/proj_out: Original Weight Key vllm-omni Expected Key Notes vace_blocks.0.before_proj.* vace_blocks.0.proj_in.* Only block 0 has before_proj vace_blocks.N.after_proj.* vace_blocks.N.proj_out.* All VACE blocks Issue 4: Text Encoder and VAE in .pth Format The original model ships models_t5_umt5-xxl-enc-bf16.pth and Wan2.1_VAE.pth as Python pickle files (.pth), not safetensors. vllm-omni's Wan22Pipeline.init loads these via HuggingFace from_pretrained(), which expects diffusers-format subdirectories with safetensors/bin files. The .pth files cannot be loaded through this path. Proposed Solutions

1. Add model_type: "vace" detection in OmniDiffusionConfig.enrich_config() to auto-map to WanVACEPipeline, similar to how "s2v" maps to WanS2VPipeline.
2. Add weight key remapping in WanTransformer3DModel.load_weights() to handle the original checkpoint's naming convention (self_attn → attn1, cross_attn → attn2, before_proj → proj_in, after_proj → proj_out, text_embedding → condition_embedder.text_embedder, etc.).
3. Support flexible subfolder names for the two-stage transformer layout (e.g., detect high_noise_model/ and low_noise_model/ as alternatives to transformer/ and transformer_2/).
4. Support .pth loading for text encoder and VAE, or provide a lightweight conversion script to turn .pth → safetensors without modifying the transformer weights. Environment

• vllm-omni: 0.20.1.dev2
• Model: Wan2.1-VACE-14B (original release, non-diffusers format)
• Config: VaceWanModel, dim=5120, num_heads=40, num_layers=40, vace_layers=0,5,10,15,20,25,30,35, vace_in_dim=96

Use case and motivation

Used for generating higher quality videos for training purposes.

Before submitting a new issue...

• Make sure you already searched for relevant issues, and asked the chatbot living at the bottom right corner of the documentation page, which can answer lots of frequently asked questions.