[Bugfix] Fix TI2V-5B weight loading by loading transformer config from model (#633)

Signed-off-by: linyueqian <linyueqian@outlook.com>

vllm_omni/diffusion/models/wan2_2/pipeline_wan2_2.py

@@ -181,16 +181,30 @@ class Wan22Pipeline(nn.Module):
 
         # Read model_index.json to detect expand_timesteps mode (for TI2V-5B)
         self.expand_timesteps = False
+        self.has_transformer_2 = False
         if local_files_only:
             model_index_path = os.path.join(model, "model_index.json")
             if os.path.exists(model_index_path):
                 with open(model_index_path) as f:
                     model_index = json.load(f)
                     self.expand_timesteps = model_index.get("expand_timesteps", False)
+            # Check if this is a two-stage model (MoE with transformer_2)
+            transformer_2_path = os.path.join(model, "transformer_2")
+            self.has_transformer_2 = os.path.exists(transformer_2_path)
+        else:
+            # For remote models, download and read model_index.json
+            try:
+                from huggingface_hub import hf_hub_download
 
-        # Check if this is a two-stage model (MoE with transformer_2)
-        transformer_2_path = os.path.join(model, "transformer_2") if local_files_only else None
-        self.has_transformer_2 = transformer_2_path is not None and os.path.exists(transformer_2_path)
+                model_index_path = hf_hub_download(repo_id=model, filename="model_index.json")
+                with open(model_index_path) as f:
+                    model_index = json.load(f)
+                    self.expand_timesteps = model_index.get("expand_timesteps", False)
+                    # Check transformer_2 from model_index
+                    transformer_2_info = model_index.get("transformer_2", [None, None])
+                    self.has_transformer_2 = transformer_2_info[0] is not None
+            except Exception:
+                pass
 
         # Set up weights sources for transformer(s)
         self.weights_sources = [
@@ -222,17 +236,13 @@ class Wan22Pipeline(nn.Module):
         ).to(self.device)
 
         # Initialize transformers with correct config (weights loaded via load_weights)
-        if local_files_only:
-            transformer_config = load_transformer_config(model, "transformer")
-            self.transformer = create_transformer_from_config(transformer_config)
-            if self.has_transformer_2:
-                transformer_2_config = load_transformer_config(model, "transformer_2")
-                self.transformer_2 = create_transformer_from_config(transformer_2_config)
-            else:
-                self.transformer_2 = None
+        transformer_config = load_transformer_config(model, "transformer", local_files_only)
+        self.transformer = create_transformer_from_config(transformer_config)
+        if self.has_transformer_2:
+            transformer_2_config = load_transformer_config(model, "transformer_2", local_files_only)
+            self.transformer_2 = create_transformer_from_config(transformer_2_config)
         else:
-            self.transformer = WanTransformer3DModel()
-            self.transformer_2 = WanTransformer3DModel() if self.has_transformer_2 else None
+            self.transformer_2 = None
 
         self.scheduler = FlowMatchEulerDiscreteScheduler.from_pretrained(
             model, subfolder="scheduler", local_files_only=local_files_only

vllm_omni/diffusion/models/wan2_2/pipeline_wan2_2_ti2v.py

@@ -31,8 +31,11 @@ from vllm.model_executor.models.utils import AutoWeightsLoader
 from vllm_omni.diffusion.data import DiffusionOutput, OmniDiffusionConfig
 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.wan2_2.pipeline_wan2_2 import retrieve_latents
-from vllm_omni.diffusion.models.wan2_2.wan2_2_transformer import WanTransformer3DModel
+from vllm_omni.diffusion.models.wan2_2.pipeline_wan2_2 import (
+    create_transformer_from_config,
+    load_transformer_config,
+    retrieve_latents,
+)
 from vllm_omni.diffusion.request import OmniDiffusionRequest
 
 
@@ -149,7 +152,9 @@ class Wan22TI2VPipeline(nn.Module):
         ).to(self.device)
 
         # Single transformer (TI2V uses dense 5B model, not MoE)
-        self.transformer = WanTransformer3DModel()
+        # Load config from model to get correct dimensions
+        transformer_config = load_transformer_config(model, "transformer", local_files_only)
+        self.transformer = create_transformer_from_config(transformer_config)
 
         # Scheduler
         self.scheduler = FlowMatchEulerDiscreteScheduler.from_pretrained(