[Quantization] Add FLUX.2-dev Mistral TE online FP8 support (#5136)

Signed-off-by: diyapeng <50388704+diyapeng@users.noreply.github.com>
Co-authored-by: diyapeng <50388704+diyapeng@users.noreply.github.com>

tests/diffusion/models/flux2/test_pipeline_flux2_quantization.py

@@ -0,0 +1,50 @@
+from unittest.mock import Mock, patch
+
+import pytest
+import torch
+
+from vllm_omni.diffusion.models.flux2.pipeline_flux2 import (
+    Flux2Pipeline,
+    _resolve_component_quant_config,
+)
+
+pytestmark = [pytest.mark.core_model, pytest.mark.cpu]
+
+
+def test_resolve_component_quant_config_routes_components():
+    text_encoder_config = object()
+    component_config = Mock()
+    component_config.resolve.side_effect = {
+        "text_encoder": text_encoder_config,
+        "transformer": None,
+    }.get
+
+    assert _resolve_component_quant_config(component_config, "text_encoder") is text_encoder_config
+    assert _resolve_component_quant_config(component_config, "transformer") is None
+
+
+def test_resolve_component_quant_config_preserves_global_config():
+    global_config = object()
+    assert _resolve_component_quant_config(global_config, "transformer") is global_config
+    assert _resolve_component_quant_config(None, "text_encoder") is None
+
+
+def test_offload_transformer_before_text_encoder_weights():
+    pipeline = object.__new__(Flux2Pipeline)
+    pipeline.transformer = Mock()
+    weights = [
+        ("transformer.layer.weight", torch.empty(1)),
+        ("text_encoder.layer.weight", torch.empty(1)),
+        ("text_encoder.layer.bias", torch.empty(1)),
+    ]
+
+    iterator = iter(pipeline._offload_transformer_before_text_encoder(weights))
+    assert next(iterator)[0] == "transformer.layer.weight"
+    pipeline.transformer.to.assert_not_called()
+
+    with patch("vllm_omni.diffusion.models.flux2.pipeline_flux2.current_omni_platform.empty_cache") as empty_cache:
+        assert next(iterator)[0] == "text_encoder.layer.weight"
+        assert next(iterator)[0] == "text_encoder.layer.bias"
+
+    pipeline.transformer.to.assert_called_once_with("cpu")
+    empty_cache.assert_called_once_with()