Fix FP8 Z-Image quality gate

Signed-off-by: WeiQing Chen <david6666666@users.noreply.github.com>

tests/diffusion/quantization/artifacts/fp8_z_image_quality/README.md

@@ -13,16 +13,28 @@ VLLM_OMNI_QUALITY_OUTPUT_DIR=$PWD/tests/diffusion/quantization/artifacts/fp8_z_i
 Quantized config:
 
 ```python
-{"method": "fp8", "ignored_layers": ["img_mlp"]}
+{
+    "transformer": {
+        "method": "fp8",
+        "ignored_layers": [
+            "img_mlp",
+            "layers.15..29.{attention.to_qkv,attention.to_out.0,feed_forward.w13,feed_forward.w2}",
+        ],
+    },
+    "default": None,
+}
 ```
 
-Result: failed with the current `max_lpips=0.15` threshold.
+Result: passed with `max_lpips=0.10`.
 
 | Metric | Value |
 |--------|-------|
-| LPIPS | 0.250578 |
-| PSNR | 17.042153 dB |
-| MAE | 0.076899 |
+| LPIPS | 0.096935 |
+| PSNR | 21.893253 dB |
+| MAE | 0.037334 |
+
+Timing from the scan run for this config: BF16 `2.297244s`, FP8 `2.062006s`
+(`0.8976x` of BF16 generation time).
 
 ## Comparison
 

tests/diffusion/quantization/test_fp8_config.py

@@ -213,6 +213,35 @@ def test_integration_per_component():
     assert config.quantization_config.component_configs["vae"] is None
 
 
+def test_recursive_replace_linear_resolves_component_quantization(monkeypatch):
+    from vllm_omni.diffusion.data import OmniDiffusionConfig
+    from vllm_omni.diffusion.models import utils
+
+    model = nn.Module()
+    model.text_encoder = nn.Sequential(nn.Linear(2, 2))
+    model.transformer = nn.Sequential(nn.Linear(2, 2))
+    config = OmniDiffusionConfig(
+        model="test",
+        quantization_config={
+            "transformer": {"method": "fp8"},
+            "default": None,
+        },
+    )
+
+    calls = {}
+
+    def fake_replace(linear, style, quant_config, *, prefix):
+        calls[prefix] = quant_config
+        return nn.Identity()
+
+    monkeypatch.setattr(utils, "replace_linear_class", fake_replace)
+
+    utils.recursive_replace_linear(model, config)
+
+    assert calls["text_encoder.0"] is None
+    assert calls["transformer.0"].get_name() == "fp8"
+
+
 def test_transformer_config_auto_detects_modelopt_fp8():
     from vllm.model_executor.layers.quantization.modelopt import ModelOptFp8Config
 

tests/diffusion/quantization/test_quantization_quality.py

@@ -118,7 +118,25 @@ QUALITY_CONFIGS = [
     QualityTestConfig(
         id="fp8_z_image",
         model="Tongyi-MAI/Z-Image-Turbo",
-        quantization={"method": "fp8", "ignored_layers": ["img_mlp"]},
+        quantization={
+            "transformer": {
+                "method": "fp8",
+                "ignored_layers": [
+                    "img_mlp",
+                    *[
+                        f"layers.{layer_id}.{suffix}"
+                        for layer_id in range(15, 30)
+                        for suffix in (
+                            "attention.to_qkv",
+                            "attention.to_out.0",
+                            "feed_forward.w13",
+                            "feed_forward.w2",
+                        )
+                    ],
+                ],
+            },
+            "default": None,
+        },
         task="t2i",
         prompt=(
             "A breathtaking twilight scene atop a floating archipelago of crystalline islands suspended in an "
@@ -141,7 +159,7 @@ QUALITY_CONFIGS = [
             "starlight waterfalls. Ethereal, melancholic, and transcendent mood—like a moment of quiet revelation "
             "at the edge of existence."
         ),
-        max_lpips=0.15,
+        max_lpips=0.10,
         num_inference_steps=20,
     ),
     QualityTestConfig(

vllm_omni/diffusion/models/utils.py

@@ -19,6 +19,7 @@ from vllm.model_executor.models.utils import maybe_prefix
 
 from vllm_omni.diffusion.data import OmniDiffusionConfig
 from vllm_omni.quantization import build_quant_config
+from vllm_omni.quantization.component_config import ComponentQuantizationConfig
 
 if TYPE_CHECKING:
     from transformers import PretrainedConfig, PreTrainedModel
@@ -87,7 +88,17 @@ def recursive_replace_linear(model: nn.Module, od_config: OmniDiffusionConfig):
             # Replace modules as needed
             if isinstance(child_module, nn.Linear):
                 style = "replicate"
-                new_module = replace_linear_class(child_module, style, quant_config, prefix=qual_name)
+                child_quant_config = (
+                    quant_config.resolve(qual_name)
+                    if isinstance(quant_config, ComponentQuantizationConfig)
+                    else quant_config
+                )
+                new_module = replace_linear_class(
+                    child_module,
+                    style,
+                    child_quant_config,
+                    prefix=qual_name,
+                )
             else:
                 _recursive_replace(child_module, prefix=qual_name)
             if new_module is not child_module:

vllm_omni/diffusion/models/z_image/pipeline_z_image.py

@@ -45,6 +45,7 @@ from vllm_omni.diffusion.request import OmniDiffusionRequest
 from vllm_omni.model_executor.model_loader.weight_utils import (
     download_weights_from_hf_specific,
 )
+from vllm_omni.quantization.component_config import ComponentQuantizationConfig
 
 logger = logging.get_logger(__name__)  # pylint: disable=invalid-name
 
@@ -218,7 +219,10 @@ class ZImagePipeline(nn.Module, DiffusionPipelineProfilerMixin):
 
         vae_config = DistributedAutoencoderKL.load_config(model, subfolder="vae", local_files_only=local_files_only)
         self.vae = DistributedAutoencoderKL.from_config(vae_config).to(self._execution_device)
-        self.transformer = ZImageTransformer2DModel(quant_config=od_config.quantization_config)
+        quant_config = od_config.quantization_config
+        if isinstance(quant_config, ComponentQuantizationConfig):
+            quant_config = quant_config.resolve("transformer")
+        self.transformer = ZImageTransformer2DModel(quant_config=quant_config)
         self.tokenizer = AutoTokenizer.from_pretrained(model, subfolder="tokenizer", local_files_only=local_files_only)
 
         # Note: Context parallelism is applied centrally in registry.initialize_model()