Add Qwen Image Edit 2509 pipeline support

README.md

@@ -35,7 +35,7 @@ multi-component export for pipelines.
 | **Speech-to-Text** | Whisper, FastConformer-RNNT, FunASR, Qwen3-ASR, SenseVoice |
 | **Audio** | Wav2Vec2, HuBERT, WavLM, SpeechT5 |
 | **Vision** | ViT, BEiT, DeiT, DINOv2, Swin, CLIP, SigLIP |
-| **Diffusion** | Stable Diffusion (UNet + VAE + ControlNet), Flux, SD3, DiT, QwenImage, HunyuanDiT, CogVideoX |
+| **Diffusion** | Stable Diffusion (UNet + VAE + ControlNet), Flux, SD3, DiT, QwenImage / Qwen-Image-Edit-2509, HunyuanDiT, CogVideoX |
 | **Adapters** | T2I-Adapter, IP-Adapter |
 
 Supports **290+ Transformers model types** and **10 Diffusers component types**

docs/model-catalog.md

@@ -243,8 +243,12 @@ Supported component classes include:
 - `ControlNetModel` — ControlNet conditioning
 - `CogVideoXTransformer3DModel` — CogVideoX
 - `VideoAutoencoderModel` — Video VAE
-- `QwenImageTransformer2DModel` — Qwen image generation
-- `AutoencoderKLQwenImageModel` — Qwen image VAE
+- `QwenImageTransformer2DModel` — Qwen image generation and packed-token
+  Qwen-Image-Edit-2509 denoising with source-image conditioning, masks, and 3D RoPE
+- `AutoencoderKLQwenImageModel` — Qwen image VAE, including edit-pipeline latent
+  normalization
+- `Qwen2_5_VLForConditionalGeneration` — image-aware prompt encoder used by
+  Qwen-Image-Edit-2509
 
 ```python
 from mobius import build

src/mobius/__main__.py

@@ -262,10 +262,21 @@ def _resolve_static_cache_task(model_type: str) -> ModelTask:
             print(
                 f"Detected diffusers pipeline: {pipeline_index.get('_class_name', 'Unknown')}"
             )
+            pipeline_components = None
+            if component_filter:
+                roots = [
+                    name
+                    for name in pipeline_index
+                    if not name.startswith("_")
+                    and (component_filter == name or component_filter.startswith(f"{name}_"))
+                ]
+                pipeline_components = {max(roots, key=len)} if roots else {component_filter}
             pkg = build_diffusers_pipeline(
                 args.model,
                 dtype=dtype_override,
                 load_weights=load_weights,
+                components=pipeline_components,
+                execution_provider=execution_provider,
             )
             _save_package(pkg, output_dir, args, optimize, component_filter)
             return

src/mobius/_builder.py

@@ -536,6 +536,7 @@ def build(
                 model_id,
                 dtype=dtype,
                 load_weights=load_weights,
+                execution_provider=execution_provider,
             )
 
     model_type = hf_config.model_type

src/mobius/_diffusers_builder.py

@@ -46,6 +46,7 @@ def _init_diffusers_class_map() -> None:
         CLIPTextConfig,
         CogVideoXConfig,
         QwenImageConfig,
+        QwenImageTextEncoderConfig,
         QwenImageVAEConfig,
         UNet2DConfig,
         VAEConfig,
@@ -64,6 +65,7 @@ def _init_diffusers_class_map() -> None:
     from mobius.models.hunyuan_dit import HunyuanDiT2DModel, HunyuanDiTConfig
     from mobius.models.qwen_image import QwenImageTransformer2DModel
     from mobius.models.qwen_image_vae import AutoencoderKLQwenImageModel
+    from mobius.models.qwen_vl import Qwen25VLCausalLMModel
     from mobius.models.unet import UNet2DConditionModel
     from mobius.models.vae import AutoencoderKLModel
     from mobius.models.video_vae import VideoAutoencoderModel, VideoVAEConfig
@@ -86,7 +88,12 @@ def _init_diffusers_class_map() -> None:
             "QwenImageTransformer2DModel": (
                 QwenImageTransformer2DModel,
                 QwenImageConfig,
-                "denoising",
+                "qwen-image-denoising",
+            ),
+            "Qwen2_5_VLForConditionalGeneration": (
+                Qwen25VLCausalLMModel,
+                QwenImageTextEncoderConfig,
+                "qwen-image-text-encoding",
             ),
             "AutoencoderKL": (AutoencoderKLModel, VAEConfig, "vae"),
             "AutoencoderKLQwenImage": (
@@ -203,6 +210,19 @@ def _load_diffusers_component_config(model_id: str, component_name: str) -> dict
         return json.load(f)
 
 
+def _load_optional_diffusers_json(model_id: str, filename: str) -> dict:
+    """Load optional non-neural pipeline metadata without failing the build."""
+    from huggingface_hub import hf_hub_download
+    from huggingface_hub.utils import EntryNotFoundError
+
+    try:
+        path = hf_hub_download(repo_id=model_id, filename=filename)
+    except EntryNotFoundError:
+        return {}
+    with open(path) as f:
+        return json.load(f)
+
+
 def _prepare_unet_loras(unet_loras: dict) -> tuple[tuple, dict]:
     """Load each UNet LoRA ``.safetensors``; return baked-adapter specs + merged weights.
 
@@ -236,6 +256,8 @@ def build_diffusers_pipeline(
     dtype: str | ir.DataType | None = None,
     load_weights: bool = True,
     unet_loras: dict | None = None,
+    components: set[str] | None = None,
+    execution_provider: str = "default",
 ) -> ModelPackage:
     """Build ONNX models for all supported components in a diffusers pipeline.
 
@@ -255,6 +277,9 @@ def build_diffusers_pipeline(
             inferred from the file); at inference a ``lora_gate.{name}`` scalar
             input switches/blends it. Requires ``load_weights=True`` to apply the
             adapter weights.
+        components: Optional component-name allowlist. Non-neural pipeline metadata
+            is still retained so a single-component export preserves its contract.
+        execution_provider: Target execution provider for EP-aware graph optimization.
 
     Returns:
         A :class:`ModelPackage` containing the built component model(s).
@@ -275,10 +300,14 @@ def build_diffusers_pipeline(
         dtype = resolve_dtype(dtype)
 
     package = ModelPackage({})
+    component_configs: dict[str, dict] = {}
+    pipeline_class = str(pipeline_index.get("_class_name", "DiffusionPipeline"))
 
     for component_name, component_info in pipeline_index.items():
         if component_name.startswith("_"):
             continue
+        if components is not None and component_name not in components:
+            continue
         if not isinstance(component_info, list) or len(component_info) != 2:
             continue
 
@@ -300,6 +329,7 @@ def build_diffusers_pipeline(
         )
 
         component_config_dict = _load_diffusers_component_config(model_id, component_name)
+        component_configs[component_name] = component_config_dict
         config = config_class.from_diffusers(component_config_dict)
 
         if dtype is not None and hasattr(config, "dtype"):
@@ -318,16 +348,29 @@ def build_diffusers_pipeline(
 
         model_module = module_class(config)
 
-        sub_pkg = build_from_module(model_module, config, task_name)
+        if (
+            pipeline_class == "QwenImageEditPlusPipeline"
+            and class_name == "AutoencoderKLQwenImage"
+        ):
+            task_name = "qwen-image-edit-vae"
+        sub_pkg = build_from_module(
+            model_module,
+            config,
+            task_name,
+            execution_provider=execution_provider,
+        )
 
         # Flatten sub-package into the top-level package
         if len(sub_pkg) == 1 and "model" in sub_pkg:
             sub_pkg["model"].graph.name = f"{model_id}/{component_name}"
             package[component_name] = sub_pkg["model"]
         else:
             for sub_name, sub_model in sub_pkg.items():
-                sub_model.graph.name = f"{model_id}/{component_name}_{sub_name}"
-                package[f"{component_name}_{sub_name}"] = sub_model
+                package_name = (
+                    component_name if sub_name == "model" else f"{component_name}_{sub_name}"
+                )
+                sub_model.graph.name = f"{model_id}/{package_name}"
+                package[package_name] = sub_model
 
         if load_weights:
             state_dict = _download_diffusers_component_weights(model_id, component_name)
@@ -344,4 +387,24 @@ def build_diffusers_pipeline(
             f"Supported diffusers classes: {sorted(_DIFFUSERS_CLASS_MAP)}."
         )
 
+    from mobius._diffusers_configs import DiffusersPipelineConfig
+
+    package.config = DiffusersPipelineConfig(
+        source_model_id=model_id,
+        pipeline_class=pipeline_class,
+        component_configs=component_configs,
+        scheduler_config=(
+            _load_optional_diffusers_json(model_id, "scheduler/scheduler_config.json")
+            if "scheduler" in pipeline_index
+            else {}
+        ),
+        processor_config=(
+            _load_optional_diffusers_json(model_id, "processor/preprocessor_config.json")
+            if "processor" in pipeline_index
+            else {}
+        ),
+        model_type=(
+            "qwen_image_edit" if pipeline_class == "QwenImageEditPlusPipeline" else "diffusers"
+        ),
+    )
     return package

src/mobius/_diffusers_builder_test.py

@@ -53,6 +53,7 @@ def test_populates_expected_classes(self):
             "FluxTransformer2DModel",
             "SD3Transformer2DModel",
             "QwenImageTransformer2DModel",
+            "Qwen2_5_VLForConditionalGeneration",
             "UNet2DConditionModel",
             "CLIPTextModel",
             "AutoencoderKL",
@@ -81,6 +82,8 @@ def test_task_names_are_valid(self):
             "denoising",
             "vae",
             "qwen-image-vae",
+            "qwen-image-denoising",
+            "qwen-image-text-encoding",
             "video-denoising",
             "feature-extraction",
         }
@@ -346,7 +349,7 @@ def test_multiple_components_built(
         )
         mock_load_config.return_value = {}
 
-        def fake_build(module, config, task_name):
+        def fake_build(module, config, task_name, **kwargs):
             graph = ir.Graph([], [], nodes=[], name="g")
             return ModelPackage({"model": ir.Model(graph, ir_version=10)})
 
@@ -400,6 +403,69 @@ def test_dtype_ir_datatype_passthrough(
         # Verify build_from_module was called (ir.DataType accepted without error)
         mock_build_from_module.assert_called_once()
 
+    @patch("mobius._diffusers_builder.build_from_module")
+    @patch("mobius._diffusers_builder._load_diffusers_component_config")
+    @patch("mobius._diffusers_builder._load_diffusers_pipeline_index")
+    def test_qwen_edit_uses_normalized_vae_task(
+        self,
+        mock_load_index,
+        mock_load_config,
+        mock_build_from_module,
+    ):
+        mock_load_index.return_value = {
+            "_class_name": "QwenImageEditPlusPipeline",
+            "vae": ["diffusers", "AutoencoderKLQwenImage"],
+        }
+        mock_load_config.return_value = {
+            "base_dim": 8,
+            "z_dim": 4,
+            "dim_mult": [1, 2],
+            "num_res_blocks": 1,
+            "temperal_downsample": [False],
+            "latents_mean": [0.0] * 4,
+            "latents_std": [1.0] * 4,
+        }
+        graph = ir.Graph([], [], nodes=[], name="vae")
+        mock_build_from_module.return_value = ModelPackage(
+            {"model": ir.Model(graph, ir_version=10)}
+        )
+
+        result = build_diffusers_pipeline("fake/qwen-edit", load_weights=False)
+
+        assert "vae" in result
+        assert mock_build_from_module.call_args.args[2] == "qwen-image-edit-vae"
+        assert result.config.model_type == "qwen_image_edit"
+
+    @patch("mobius._diffusers_builder.build_from_module")
+    @patch("mobius._diffusers_builder._load_diffusers_component_config")
+    @patch("mobius._diffusers_builder._load_diffusers_pipeline_index")
+    def test_component_allowlist_avoids_building_other_components(
+        self,
+        mock_load_index,
+        mock_load_config,
+        mock_build_from_module,
+    ):
+        mock_load_index.return_value = _fake_pipeline_index(
+            {
+                "transformer": ["diffusers", "FluxTransformer2DModel"],
+                "vae": ["diffusers", "AutoencoderKL"],
+            }
+        )
+        mock_load_config.return_value = {}
+        graph = ir.Graph([], [], nodes=[], name="transformer")
+        mock_build_from_module.return_value = ModelPackage(
+            {"model": ir.Model(graph, ir_version=10)}
+        )
+
+        result = build_diffusers_pipeline(
+            "fake/filtered",
+            load_weights=False,
+            components={"transformer"},
+        )
+
+        assert set(result) == {"transformer"}
+        mock_load_config.assert_called_once_with("fake/filtered", "transformer")
+
 
 # ── build_diffusers_pipeline weight loading ──────────────────────────────
 
@@ -504,7 +570,7 @@ def test_preprocess_weights_called_when_available(
 
         # The module class will have preprocess_weights set by AutoencoderKLModel
         # We patch it at the module instance level via build_from_module's first arg
-        def capture_build(module, config, task_name):
+        def capture_build(module, config, task_name, **kwargs):
             module.preprocess_weights = lambda sd: processed_weights
             return ModelPackage({"model": model})
 

src/mobius/_diffusers_configs.py

@@ -8,6 +8,8 @@
 import dataclasses
 from typing import TYPE_CHECKING
 
+import onnx_ir as ir
+
 if TYPE_CHECKING:
     from mobius._configs import ArchitectureConfig
 
@@ -55,6 +57,37 @@ def from_diffusers(cls, config: dict) -> ArchitectureConfig:
         )
 
 
+class QwenImageTextEncoderConfig:
+    """Adapter for the Qwen2.5-VL prompt encoder bundled with Qwen Image Edit."""
+
+    @classmethod
+    def from_diffusers(cls, config: dict) -> ArchitectureConfig:
+        """Build the native Mobius Qwen2.5-VL configuration tree."""
+        import transformers
+
+        from mobius._configs import ArchitectureConfig
+
+        if hasattr(config, "to_dict"):
+            config = config.to_dict()
+        fields = dict(config)
+        model_type = fields.pop("model_type", "qwen2_5_vl")
+        hf_config = transformers.AutoConfig.for_model(model_type, **fields)
+        text_config = hf_config.text_config if hasattr(hf_config, "text_config") else hf_config
+        return ArchitectureConfig.from_transformers(text_config, parent_config=hf_config)
+
+
+@dataclasses.dataclass
+class DiffusersPipelineConfig:
+    """Non-neural diffusers pipeline metadata retained on a ModelPackage."""
+
+    source_model_id: str
+    pipeline_class: str
+    component_configs: dict[str, dict]
+    scheduler_config: dict
+    processor_config: dict
+    model_type: str = "diffusers"
+
+
 @dataclasses.dataclass
 class VAEConfig:
     """Configuration for AutoencoderKL (VAE) models."""
@@ -229,6 +262,7 @@ class QwenImageConfig:
     guidance_embeds: bool = False
     axes_dims_rope: tuple[int, ...] = (16, 56, 56)
     norm_eps: float = 1e-6
+    dtype: ir.DataType = ir.DataType.FLOAT
     # cross_attention_dim is used by DenoisingTask for encoder_hidden_states shape
     cross_attention_dim: int = 3584
 
@@ -262,6 +296,9 @@ class QwenImageVAEConfig:
     attn_scales: tuple[float, ...] = ()
     temperal_downsample: tuple[bool, ...] = (False, True, True)
     dropout: float = 0.0
+    latents_mean: tuple[float, ...] | None = None
+    latents_std: tuple[float, ...] | None = None
+    dtype: ir.DataType = ir.DataType.FLOAT
 
     @classmethod
     def from_diffusers(cls, config: dict) -> QwenImageVAEConfig:
@@ -276,4 +313,12 @@ def from_diffusers(cls, config: dict) -> QwenImageVAEConfig:
             attn_scales=tuple(config.get("attn_scales", [])),
             temperal_downsample=tuple(config.get("temperal_downsample", [False, True, True])),
             dropout=config.get("dropout", 0.0),
+            latents_mean=(
+                tuple(config["latents_mean"])
+                if config.get("latents_mean") is not None
+                else None
+            ),
+            latents_std=(
+                tuple(config["latents_std"]) if config.get("latents_std") is not None else None
+            ),
         )