Add Qwen Image Edit 2509 pipeline support - #465
Open
justinchuby wants to merge 3 commits into
Open
Conversation
Implement the Qwen-Image-Edit-2509 packed-latent transformer contract, Qwen2.5-VL prompt encoder split, normalized Qwen VAE path, FlowMatch scheduler metadata, processor assets, and deterministic diffusion goldens. Co-authored-by: Copilot App <223556219+Copilot@users.noreply.github.com> Signed-off-by: Justin Chu <justinchu@microsoft.com>
![github-advanced-security[bot]](https://avatars.githubusercontent.com/in/57789?s=60&v=4){kind=small}
Performance Comparison
|
🏗️ Architecture Diff
No architecture changes detected. ✅ Legend: ⚪ No change · 🔵 Minor (attrs/inits) · 🟡 Moderate (nodes added/removed) · 🔴 Major (interface changed) |
Contributor
There was a problem hiding this comment.
Pull request overview
Adds end-to-end support for the Qwen/Qwen-Image-Edit-2509 diffusers pipeline in mobius, including model/task wiring for prompt encoding + denoising + edit-VAE, plus onnx-genai export metadata so the runtime can reproduce the pipeline’s non-neural contract (latent packing, RoPE inputs, scheduler shifting, etc.).
Changes:
- Introduces Qwen Image Edit–specific tasks (prompt text-encoder split, packed-token denoiser, normalized-latent VAE) and updates Qwen Image/Qwen VL model codepaths to match diffusers behavior.
- Extends diffusers pipeline building and onnx-genai auto-export to retain/emit pipeline metadata (scheduler + processor assets) and support filtered/single-component exports.
- Adds unit/integration coverage and deterministic L4/L5 diffusion-stage goldens for the new pipeline.
Reviewed changes
Copilot reviewed 24 out of 24 changed files in this pull request and generated 2 comments.
Show a summary per file
| File | Description |
|---|---|
| tests/integration_test.py | Uses TemporaryDirectory for ORT session model saving to improve portability/cleanup. |
| tests/build_graph_test.py | Adds graph-build coverage for Qwen Image Edit text-encoder split and updates Qwen Image denoising task expectations. |
| testdata/golden/diffusion/qwen-image-edit-2509.json | Adds deterministic diffusion-stage golden outputs for the Qwen Image Edit 2509 pipeline. |
| testdata/cases/diffusion/qwen-image-edit-2509.yaml | Adds an L4/L5 golden-case descriptor for the official Qwen Image Edit 2509 pipeline. |
| src/mobius/tasks/_qwen_image.py | New packed-token denoising task for Qwen Image with external RoPE + mask + target slicing. |
| src/mobius/tasks/_qwen_image_vae.py | Adds dtype plumbing for VAE graphs and introduces an edit-VAE task that embeds latent normalization. |
| src/mobius/tasks/_qwen_image_text_encoder.py | New task that builds the Qwen2.5-VL prompt encoder split needed by image edit. |
| src/mobius/tasks/init.py | Exposes new Qwen Image tasks and registers their task names. |
| src/mobius/models/qwen_vl.py | Adds return_hidden_states option to support prompt-encoder hidden-state extraction. |
| src/mobius/models/qwen_image.py | Implements Qwen Image Edit denoiser changes: packed tokens, 3D RoPE inputs, masks, timestep embedding changes, output slicing, fp16 clipping. |
| src/mobius/models/qwen_image_vae.py | Adjusts RMSNorm broadcasting for attention paths and preserves diffusers weight names via passthrough preprocess. |
| src/mobius/models/qwen_image_test.py | Adds numerical parity tests vs diffusers, CUDA/precision coverage, and deterministic L4/L5 diffusion-stage golden replay. |
| src/mobius/integrations/onnx_genai/inference_metadata.py | Extends scheduler metadata to support FlowMatch Euler fields and omits beta fields for flow-match configs. |
| src/mobius/integrations/onnx_genai/inference_metadata_test.py | Adds coverage for FlowMatch Euler scheduler field mapping and metadata emission. |
| src/mobius/integrations/onnx_genai/auto_export.py | Adds Qwen Image Edit–specific metadata augmentation, processor asset copying, and flat-path handling for single-component diffusion exports. |
| src/mobius/integrations/onnx_genai/auto_export_test.py | Tests flat-path diffusion export and Qwen Image Edit metadata + processor assets emission. |
| src/mobius/_diffusers_configs.py | Adds Qwen2.5-VL text-encoder adapter config and stores non-neural pipeline metadata on ModelPackage; adds dtype + latent stats to Qwen configs. |
| src/mobius/_diffusers_configs_test.py | Adds L2-ish config mapping coverage for official Qwen Image Edit component configs (transformer + VAE latent stats). |
| src/mobius/_diffusers_builder.py | Adds Qwen2.5-VL component support, component allowlisting, EP-aware build plumbing, normalized VAE task selection for Qwen Image Edit, and retains pipeline metadata. |
| src/mobius/_diffusers_builder_test.py | Updates class/task expectations and adds tests for normalized VAE task selection + component allowlist behavior. |
| src/mobius/_builder.py | Passes execution_provider through to diffusers pipeline building. |
| src/mobius/main.py | Adds CLI component filtering behavior for diffusers pipelines and forwards EP selection through pipeline build. |
| README.md | Documents diffusion support including Qwen-Image-Edit-2509. |
| docs/model-catalog.md | Documents supported diffusers component classes including Qwen Image Edit prompt encoder + normalized VAE behavior. |
Suppressed comments (5)
src/mobius/models/qwen_image_test.py:187
- This path uses a Windows-style "\" separator. On POSIX it produces a filename containing a backslash and writes outside the temp directory, leaving artifacts behind.
path = f"{directory}\\model.onnx"
src/mobius/models/qwen_image_test.py:192
- Profiling output prefix uses a Windows-style "\" separator; on POSIX this writes profiling JSON outside the temp directory (and may leave it behind).
options.profile_file_prefix = f"{directory}\\cuda_profile"
src/mobius/models/qwen_image_test.py:371
- This path uses a Windows-style "\" separator. On POSIX it writes the model outside the temp directory (backslash becomes part of the filename), which can leak files between tests.
path = f"{directory}\\model.onnx"
src/mobius/models/qwen_image_test.py:486
- This path uses a Windows-style "\" separator. On POSIX it writes the model outside the temp directory, which can leak artifacts and break test isolation.
path = f"{directory}\\model.onnx"
src/mobius/tasks/_qwen_image_vae.py:111
- Same validation issue as the encoder: if
latents_mean/latents_stdlengths don’t matchz_dim, the decoder graph can be built but will fail at runtime.
if config.latents_mean is None or config.latents_std is None:
raise ValueError(
"Qwen Image Edit VAE config requires latents_mean and latents_std"
)
💡 Add a code-review agent skill or configure MCP servers for context-aware, tailored reviews. Learn more in the docs.
| ).sample[:, :4] | ||
|
|
||
| with tempfile.TemporaryDirectory() as directory: | ||
| path = f"{directory}\\model.onnx" |
Member
Author
There was a problem hiding this comment.
Fixed in 4cef061: all temporary model and profiling paths now use os.path.join, preserving TemporaryDirectory isolation on POSIX and Windows. Direct Ruff format/check and the Qwen Image test module pass.
Comment on lines
+83
to
+86
| if config.latents_mean is None or config.latents_std is None: | ||
| raise ValueError( | ||
| "Qwen Image Edit VAE config requires latents_mean and latents_std" | ||
| ) |
Member
Author
There was a problem hiding this comment.
Fixed in 4cef061: QwenImageEditVAETask now validates that both latent-stat arrays exist and have exactly z_dim entries before either graph is constructed, with a regression test for the mismatch case.
Use portable temporary paths in diffusion tests and validate Qwen VAE latent statistics before graph construction. Co-authored-by: Copilot App <223556219+Copilot@users.noreply.github.com> Signed-off-by: Justin Chu <justinchu@microsoft.com>
Generate reduced L4 and L5 references from an independent diffusers VAE encode, packed transformer FlowMatch loop, and VAE decode, then replay the full Mobius ONNX chain against those outputs. Co-authored-by: Copilot App <223556219+Copilot@users.noreply.github.com> Signed-off-by: Justin Chu <justinchu@microsoft.com>
justinchuby
commented
Aug 8, 2026
|
|
||
| sub_pkg = build_from_module(model_module, config, task_name) | ||
| if ( | ||
| pipeline_class == "QwenImageEditPlusPipeline" |
Member
Author
There was a problem hiding this comment.
This doesn't look general enough
| else {} | ||
| ), | ||
| model_type=( | ||
| "qwen_image_edit" if pipeline_class == "QwenImageEditPlusPipeline" else "diffusers" |
| self.txt_norm = _RMSNorm(config.joint_attention_dim, eps=config.norm_eps) | ||
| self.txt_in = _Linear(config.joint_attention_dim, hidden_size) | ||
| self.time_text_embed = _TimestepEmbedding(hidden_size, hidden_size) | ||
| self.time_text_embed = _TimestepEmbedding(256, hidden_size) |
This file contains hidden or bidirectional Unicode text that may be interpreted or compiled differently than what appears below. To review, open the file in an editor that reveals hidden Unicode characters.
Learn more about bidirectional Unicode characters
Sign up for free
to join this conversation on GitHub.
Already have an account?
Sign in to comment
Add this suggestion to a batch that can be applied as a single commit.This suggestion is invalid because no changes were made to the code.Suggestions cannot be applied while the pull request is closed.Suggestions cannot be applied while viewing a subset of changes.Only one suggestion per line can be applied in a batch.Add this suggestion to a batch that can be applied as a single commit.Applying suggestions on deleted lines is not supported.You must change the existing code in this line in order to create a valid suggestion.Outdated suggestions cannot be applied.This suggestion has been applied or marked resolved.Suggestions cannot be applied from pending reviews.Suggestions cannot be applied on multi-line comments.Suggestions cannot be applied while the pull request is queued to merge.Suggestion cannot be applied right now. Please check back later.
Summary
QwenImageEditPlusPipelinedetection forQwen/Qwen-Image-Edit-2509, including its Qwen2.5-VL image-aware prompt encoder, Qwen Image transformer, normalized 3D VAE, processor assets, and FlowMatch Euler scheduler configurationValidation
lintrunner -a— full repository lint clean after the golden strengtheningQwen/Qwen-Image-Edit-2509fp16 CUDA component CLI export with--runtime onnx-genai— succeeded and emittedmodel.onnx, FlowMatch metadata, tokenizer, Qwen2-VL image/video processor configs, and chat templatertol=atol=1e-2os.path.joinand added strictlatents_mean/latents_stdlength validation againstz_dimGolden provenance and semantics
The committed golden is a genuine independently generated reduced Qwen Image Edit diffusion pipeline reference, but it is not generated with the official 20.4B checkpoint weights and is not an official-checkpoint perceptual output.
Reference generation uses diffusers 0.39.0 with a seeded one-block
QwenImageTransformer2DModeland seeded reducedAutoencoderKLQwenImagewhose weights are copied into the separate Mobius ONNX graphs. The source is the real nonzerotestdata/pipeline-cat-chonk.jpeg:FlowMatchEulerDiscreteSchedulerstepsThe test then independently runs the Mobius chain: ONNX VAE encoder -> separate NumPy 2x2 packing -> ONNX denoiser plus a separate FlowMatch scheduler -> NumPy unpacking -> ONNX VAE decoder.
l4_noise_pred): all 256 target-token flow-prediction values from the first actual diffusers transformer call. Both the diffusers reference and ONNX output are compared elementwise to the committed data.l5_final_latents,l5_final_image): all 256 target latent-token values after three actual diffusers transformer/scheduler steps and all 768 clipped decoded image values from the final reduced diffusers VAE decode. Mobius final latent tokens and decoded image are compared elementwise.This proves real-image VAE conditioning, packing, transformer invocation, multi-step scheduling, and decode semantics without fake token-generation data. It does not claim official-weight image quality or official-checkpoint end-to-end parity.
Waivers