fix(diffusion): honor local checkpoint dirs and task-derived model id in WAN inference

examples/models/wan/README.md

@@ -186,6 +186,16 @@ For more details, see the recipe in `src/megatron/bridge/diffusion/recipes/wan/w
 
 The script [inference_wan.py](inference_wan.py) runs text-to-video generation with a Megatron-format WAN checkpoint. Set `--checkpoint_step` to use a specific checkpoint iteration, `--sizes` and `--frame_nums` to specify video shape, and `--sample_steps` (default 50) for the number of noise diffusion steps.
 
+**Prerequisites**: Inference and video saving rely on a few packages that are **not** shipped in the container (the codecs `imageio`, `imageio-ffmpeg`, and `av` carry CVEs and are excluded; `easydict` is an example-only import). Install them into the active environment before running:
+
+```bash
+bash scripts/install_diffusion_deps.sh
+# or directly:
+uv pip install --no-config easydict imageio imageio-ffmpeg av
+```
+
+Without `easydict`, the script fails immediately at import with `ModuleNotFoundError: No module named 'easydict'`.
+
 ```bash
 uv run python -m torch.distributed.run --nproc_per_node=1 \
   examples/models/wan/inference_wan.py \

examples/models/wan/inference_wan.py

@@ -44,6 +44,14 @@
     },
 }
 
+# Maps each task to its Hugging Face diffusers model id, used to load the non-DiT
+# components (text encoder, tokenizer, VAE, scheduler) when no local checkpoint
+# directory is supplied via --t5_checkpoint_dir / --vae_checkpoint_dir.
+TASK_TO_MODEL_ID = {
+    "t2v-1.3B": "Wan-AI/Wan2.1-T2V-1.3B-Diffusers",
+    "t2v-14B": "Wan-AI/Wan2.1-T2V-14B-Diffusers",
+}
+
 
 def _validate_args(args):
     # Basic check
@@ -229,10 +237,11 @@ def generate(args):  # noqa: D103
         )
 
         logging.info("Creating flow inference pipeline.")
+        model_id = TASK_TO_MODEL_ID[args.task]
         pipeline = FlowInferencePipeline(
             inference_cfg=inference_cfg,
             checkpoint_dir=args.checkpoint_dir,
-            model_id="Wan-AI/Wan2.1-T2V-14B-Diffusers",
+            model_id=model_id,
             checkpoint_step=args.checkpoint_step,
             t5_checkpoint_dir=args.t5_checkpoint_dir,
             vae_checkpoint_dir=args.vae_checkpoint_dir,

scripts/install_diffusion_deps.sh

@@ -24,4 +24,4 @@
 # neutralize the install.
 set -euo pipefail
 
-uv pip install --no-config imageio imageio-ffmpeg av
+uv pip install --no-config imageio imageio-ffmpeg av easydict

src/megatron/bridge/diffusion/models/wan/flow_matching/flow_inference_pipeline.py

@@ -119,20 +119,29 @@ def __init__(
         self.param_dtype = inference_cfg.param_dtype
         self.text_len = inference_cfg.text_len
 
+        # Resolve where the non-DiT components (text encoder, tokenizer, VAE, scheduler)
+        # are loaded from. Prefer locally provided directories so inference can run
+        # fully offline (e.g. air-gapped nodes); fall back to the hub `model_id`.
+        t5_dir = t5_checkpoint_dir or model_id
+        vae_dir = vae_checkpoint_dir or model_id
+        # The scheduler config ships in the same diffusers repo as the other
+        # components; prefer any provided local dir, else fall back to model_id.
+        self.scheduler_source = t5_checkpoint_dir or vae_checkpoint_dir or model_id
+
         self.text_encoder = UMT5EncoderModel.from_pretrained(
-            model_id,
+            t5_dir,
             subfolder="text_encoder",
             torch_dtype=inference_cfg.t5_dtype,
         )
         self.tokenizer = AutoTokenizer.from_pretrained(
-            model_id,
+            t5_dir,
             subfolder="tokenizer",
         )
 
         self.vae_stride = inference_cfg.vae_stride
         self.patch_size = inference_cfg.patch_size
         self.vae = AutoencoderKLWan.from_pretrained(
-            model_id,
+            vae_dir,
             subfolder="vae",
             torch_dtype=inference_cfg.param_dtype,
         )
@@ -426,7 +435,9 @@ def noop_no_sync():
             batch_size_for_schedulers = len(noises)
             schedulers = []
             for _ in range(batch_size_for_schedulers):
-                base_sched = FlowMatchEulerDiscreteScheduler.from_pretrained(self.model_id, subfolder="scheduler")
+                base_sched = FlowMatchEulerDiscreteScheduler.from_pretrained(
+                    self.scheduler_source, subfolder="scheduler"
+                )
                 s = UniPCMultistepScheduler.from_config(base_sched.config, flow_shift=shift)
                 s.set_timesteps(sampling_steps, device=self.device)
 

tests/unit_tests/diffusion/model/wan/flow_matching/test_flow_inference_pipeline.py

@@ -50,6 +50,7 @@ def _make_pipeline(**overrides):
         text_len=512,
         sp_size=1,
         model_id="Wan-AI/Wan2.1-T2V-14B-Diffusers",
+        scheduler_source="Wan-AI/Wan2.1-T2V-14B-Diffusers",
         sample_neg_prompt="",
     )
     defaults.update(overrides)