[Bugfix] Fix diffusion hf_config fallback (#4093)

Signed-off-by: KexiongYu <yukexiong1@huawei.com>

examples/offline_inference/hunyuan_image3/end2end.py

@@ -21,6 +21,7 @@ from vllm_omni.diffusion.models.hunyuan_image3.prompt_utils import (
 )
 from vllm_omni.entrypoints.omni import Omni
 from vllm_omni.inputs.data import OmniPromptType
+from vllm_omni.platforms import current_omni_platform
 
 # task -> (sys_type, bot_task, trigger_tag)
 _TASK_PRESETS: dict[str, tuple[str, str | None, str | None]] = {
@@ -109,6 +110,32 @@ def parse_args():
     parser.add_argument("--log-stats", action="store_true", default=False)
     parser.add_argument("--init-timeout", type=int, default=300, help="Initialization timeout in seconds.")
     parser.add_argument("--enforce-eager", action="store_true", help="Disable torch.compile.")
+    parser.add_argument(
+        "--quantization",
+        type=str,
+        default=None,
+        choices=["fp8", "int8", "gguf"],
+        help="Quantization method for the transformer. "
+        "Options: 'fp8' (FP8 W8A8 on Ada/Hopper, weight-only on older GPUs), 'int8' (Int8 W8A8), 'gguf' (GGUF quantized weights). "
+        "Default: None (no quantization, uses BF16).",
+    )
+    parser.add_argument(
+        "--gguf-model",
+        type=str,
+        default=None,
+        help=("GGUF file path or HF reference for transformer weights. Required when --quantization gguf is set."),
+    )
+    parser.add_argument(
+        "--ignored-layers",
+        type=str,
+        default=None,
+        help="Comma-separated list of layer name patterns to skip quantization. "
+        "Only used when --quantization is set. "
+        "Available layers: to_qkv, to_out, add_kv_proj, to_add_out, img_mlp, txt_mlp, proj_out. "
+        "Example: --ignored-layers 'add_kv_proj,to_add_out'",
+    )
+
+    current_omni_platform.pre_register_and_update(parser)
 
     from vllm_omni.engine.arg_utils import nullify_stage_engine_defaults
 
@@ -125,6 +152,24 @@ def main():
 
     # Determine stage config
     stage_configs_path = args.stage_configs_path or _MODALITY_DEFAULT_CONFIG[args.modality]
+    # Build quantization kwargs: use quantization_config dict when
+    # ignored_layers is specified so the list flows through OmniDiffusionConfig
+    quant_kwargs: dict[str, Any] = {}
+    ignored_layers = [s.strip() for s in args.ignored_layers.split(",") if s.strip()] if args.ignored_layers else None
+    if args.quantization == "gguf":
+        if not args.gguf_model:
+            raise ValueError("--gguf-model is required when --quantization gguf is set.")
+        quant_kwargs["quantization_config"] = {
+            "method": "gguf",
+            "gguf_model": args.gguf_model,
+        }
+    elif args.quantization and ignored_layers:
+        quant_kwargs["quantization_config"] = {
+            "method": args.quantization,
+            "ignored_layers": ignored_layers,
+        }
+    elif args.quantization:
+        quant_kwargs["quantization"] = args.quantization
 
     # Build Omni
     omni_kwargs = {
@@ -134,6 +179,7 @@ def main():
         "log_stats": args.log_stats,
         "init_timeout": args.init_timeout,
         "enforce_eager": args.enforce_eager,
+        **quant_kwargs,
     }
     if args.modality in ("text2img", "img2img"):
         omni_kwargs["mode"] = "text-to-image"

vllm_omni/diffusion/models/wan2_2/patch_diffusers.py

@@ -9,10 +9,11 @@ from vllm_omni.diffusion.layers.norm import RMSNormVAE
 def patch_wan_rms_norm():
     """Patch diffusers Wan RMSNorm implementation with RMSNormVAE."""
 
-    # NOTE: iterate over a snapshot of sys.modules. `hasattr` can trigger lazy
-    # submodule imports (e.g. transformers' `_LazyModule.__getattr__`), which
-    # mutate sys.modules during iteration and raise
-    # `RuntimeError: dictionary changed size during iteration`.
-    for module_name, module in list(sys.modules.items()):
-        if hasattr(module, "WanRMS_norm"):
-            setattr(module, "WanRMS_norm", RMSNormVAE)
+    # Probe `__dict__` directly instead of `hasattr`: the latter triggers
+    # custom `__getattr__` hooks (e.g. transformers' image_processing alias
+    # modules), which both emit spurious deprecation warnings and can lazily
+    # import submodules, mutating sys.modules mid-iteration.
+    for module in list(sys.modules.values()):
+        module_dict = getattr(module, "__dict__", None)
+        if module_dict is not None and "WanRMS_norm" in module_dict:
+            module_dict["WanRMS_norm"] = RMSNormVAE

vllm_omni/diffusion/worker/diffusion_worker.py

@@ -125,9 +125,9 @@ class DiffusionWorker:
         try:
             hf_config = get_config(self.od_config.model, trust_remote_code=self.od_config.trust_remote_code)
         except ValueError:
-            hf_config = None
-            logger.info("Skipping hf_config loading for diffusion model %r", self.od_config.model_class_name)
-        hf_text_config = get_hf_text_config(hf_config) if hf_config is not None else None
+            hf_config = getattr(self.od_config, "tf_model_config", None)
+            logger.info("Using diffusion tf_model_config for model %r", self.od_config.model_class_name)
+        hf_text_config = get_hf_text_config(hf_config) if hasattr(hf_config, "get_text_config") else hf_config
         vllm_config.model_config = SimpleNamespace(
             hf_config=hf_config,
             hf_text_config=hf_text_config,