Repository navigation
Fix GGUF config mapping for scaled RoPE and hybrid attention - #458
Closed
justinchuby wants to merge 46 commits into
Closed
justinchuby wants to merge 46 commits into
justinchuby wants to merge 46 commits into
Conversation
Co-authored-by: Copilot <223556219+Copilot@users.noreply.github.com>
…pipelines Mobius builds a diffusion model's neural components (denoiser transformer, VAE) but carries no scheduler loop. onnx-genai's *iterative* pipeline supplies that loop declaratively from an `inference_metadata` document. This adds `mobius.integrations.onnx_genai` to emit that document so a Mobius-built diffusion package is directly runnable by onnx-genai. - `build_diffusion_pipeline_metadata(...)` / `write_diffusion_pipeline_metadata(...)`: produce a `kind: iterative` pipeline with the denoiser loop-carried self-edge (`noise_pred -> sample`), per-step `timestep_input`, a `scheduler_config` (DDIM: num_train_timesteps/beta_start/beta_end/ prediction_type), optional VAE (`final_only`, decodes the final latent) and text encoder (`prompt_only`, feeds conditioning), and classifier-free guidance (`guidance_scale` + `cfg_conditioning_input`) when scale != 1.0. - `SchedulerConfig` with `.from_diffusers(...)` to read a diffusers scheduler_config.json. - Runtime-agnostic: reads only plain values, no torch/diffusers state. Tests (7) cover denoiser-only, full VAE+text-encoder+CFG, guidance=1.0 (no CFG), diffusers scheduler parsing, zero-steps rejection, YAML round-trip, and — importantly — validation of the emitted document against onnx-genai's published JSON schema (schema/inference_metadata.schema.json). Verified end-to-end against onnx-genai: a Mobius-built tiny DiT denoiser + this emitted metadata runs through onnx-genai's DDIM-scheduled iterative pipeline (onnx-genai PR #35, real_dit e2e test).
…steps - SchedulerConfig gains beta_schedule (default scaled_linear, the SD schedule) and emits it; from_diffusers reads it. - build_diffusion_pipeline_metadata gains a timesteps= param (explicit per-step schedule, length must == num_inference_steps) and defaults the VAE latent input to 'latent'. - Verified end-to-end: onnx-genai runs a full text-encoder -> denoiser(DDIM+CFG) -> VAE pipeline from this emitted metadata, pixel-identical to diffusers (onnx-genai scripts/diffusion_e2e.py).
… + CLI runtime Extends the onnx-genai integration beyond diffusion to language models and wires it into the CLI, so `mobius build ... --runtime onnx-genai` emits a runnable inference_metadata.yaml for both model families. - decoder_metadata: build_decoder_metadata / decoder_metadata_from_config / write_decoder_metadata — map a Mobius BaseModelConfig/ArchitectureConfig (num_attention_heads, num_key_value_heads, head_dim, max_position_embeddings, sliding_window) onto onnx-genai's `model.attention` (multi_head/grouped_query) + `kv_cache.native_dtype`. Drops DEFAULT_INT sentinels; derives head_dim from hidden_size when absent. - auto_export.write_onnx_genai_config: dispatches on the package — diffusion (denoiser/vae) -> iterative pipeline document; else -> decoder document. - CLI: `--runtime` now accepts `onnx-genai` and writes inference_metadata.yaml. Tests (17 total): decoder GQA/MHA/sliding-window/sink, from_config derivation + DEFAULT_INT dropping, YAML round-trip, schema validation against onnx-genai's published JSON schema; plus the dispatcher for both families.
SchedulerConfig.from_diffusers now maps the diffusers _class_name to the right onnx-genai scheduler kind (DDIMScheduler->ddim, EulerDiscreteScheduler->euler) instead of always emitting ddim, and rejects ancestral/SDE and other unsupported scheduler classes with a clear error so a built package never silently runs the wrong denoise dynamics. Add load_diffusers_scheduler_config(source): best-effort read of a checkpoint's scheduler/scheduler_config.json (local dir or HF hub); unsupported/absent -> None (warn + caller falls back to DDIM), never fatal. Wire it through write_onnx_genai_config(source=...) and the CLI (--runtime onnx-genai) so a diffusion build auto-populates the scheduler from the source checkpoint. Tests: euler mapping, ancestral/unsupported rejection, local-dir load, absent->None, unsupported->None fallback, and end-to-end auto-read via write_onnx_genai_config. 25 onnx_genai integration tests pass.
… metadata Add mobius.integrations.onnx_genai.comfyui: a proof-of-concept translator that maps a ComfyUI *API-format* workflow JSON onto onnx-genai's composite iterative pipeline. The KSampler-centric txt2img graph maps directly: KSampler.steps -> num_inference_steps KSampler.cfg -> guidance_scale (CFG; 1.0 disables) KSampler.sampler_name-> scheduler kind (euler, ddim) CLIPTextEncode(+/-) -> text_encoder component + CFG cond/uncond VAEDecode -> vae component (final phase) It reuses build_diffusion_pipeline_metadata, so the output is the same schema a Mobius diffusion build emits (topology/params only; weights come from the ONNX components). Unsupported samplers (dpmpp_2m, euler_ancestral, ...) and non-txt2img graphs raise a clear ValueError rather than silently running wrong dynamics. Tests (10, incl. schema conformance against onnx-genai's published inference_metadata.schema.json): default txt2img, ddim mapping, cfg=1 disables guidance, prompt-wrapper, unsupported/missing/multiple sampler rejection, denoiser-only graph, file loader. 35 onnx_genai integration tests pass.
…xport+CLI) Turn a ComfyUI API-format workflow into a runnable onnx-genai pipeline: - comfyui.py: parse_comfyui_workflow walks the graph links to recover full run params (prompt/negative via CLIPTextEncode, seed, width/height via EmptyLatentImage, checkpoint traced through LoraLoader/etc.) plus steps/cfg/sampler; maps sampler_name->scheduler kind (euler/ddim), rejects unsupported/ancestral samplers. - checkpoint_export.py: export_checkpoint bridges safetensors->ONNX via diffusers (from_single_file for ComfyUI original-SD single files, from_pretrained for diffusers dirs / HF ids), emitting text_encoder/denoiser/vae with the exact ports onnx-genai's validated diffusion_e2e contract expects (input_ids-> last_hidden_state; sample/timestep/encoder_hidden_states->noise_pred; latent->image with 1/scaling baked in), dynamic spatial/seq axes. - convert.py: convert_comfyui_workflow reconciles the ComfyUI sampler (kind/ steps/cfg) with the checkpoint's own noise schedule (betas/num_train_timesteps), computes exact diffusers timesteps, writes inference_metadata.yaml + run.json. - CLI: 'mobius convert-comfyui <workflow.json> --checkpoint <src> -o <dir>'. Validated end-to-end: converting a real txt2img workflow against OFA-Sys/small-stable-diffusion-v0 produces a complete composite pipeline (euler+CFG, correct timesteps) + 3 ONNX components. 41 integration tests pass (translator graph-resolution, checkpoint tracing, reconciliation, schema conformance).
…eps) Euler's linspace inference timesteps are fractional (e.g. 856.29); exporting the denoiser with an int64 timestep port truncated them before the UNet time embedding, causing localized pixel error. export_checkpoint gains timestep_dtype (int64 default, float32 for fractional-timestep schedulers) and convert_comfyui_workflow selects float32 for euler. End-to-end ComfyUI->onnx-genai render vs diffusers Euler on OFA-Sys/small-stable-diffusion-v0 improved from max|diff| 0.43 -> 7.9e-3 (mean 2.5e-5); also validates onnx-genai's Euler scheduler on a real SD UNet.
ComfyUI sampler_name 'dpmpp_2m' and diffusers DPMSolverMultistepScheduler now map to onnx-genai's new dpmpp_2m scheduler kind; convert.py computes the exact DPM timesteps via diffusers (int64 timestep, unlike Euler's float32). DPM++ 2M is the default sampler in most SD/ComfyUI workflows, so this substantially widens ComfyUI coverage. Tests updated (dpmpp mapping in comfyui + from_diffusers; unsupported examples switched to LMS/dpmpp_2m_sde). 43 integration tests pass.
Add use_karras_sigmas to SchedulerConfig (emitted only when True); from_diffusers reads it; ComfyUI scheduler='karras' enables it; convert.py computes Karras timesteps via diffusers and uses a float32 denoiser timestep (Karras timesteps are fractional). Enables 'DPM++ 2M Karras' — the most popular SD sampler combo. 45 integration tests pass.
The ComfyUI translator reads KSampler.denoise; when <1.0 it computes start_step = num_steps - round(num_steps*denoise) (matching diffusers get_timesteps) and emits it in the pipeline metadata (build_diffusion_pipeline_metadata gains a start_step arg). ComfyUIWorkflow carries denoise/start_step; convert threads it through. 47 integration tests pass (denoise->start_step mapping, denoise=1 omits it).
ComfyUI sampler_name 'euler_ancestral' and diffusers EulerAncestralDiscreteScheduler now map to onnx-genai's new stochastic euler_ancestral kind; convert.py computes its timesteps (float32, fractional) via diffusers. Other ancestral/SDE samplers still rejected with a clear error. 49 integration tests pass.
Detect SDXL checkpoints (pipe.text_encoder_2) via DiffusionPipeline auto-loading and export the SDXL shape: a combined dual text encoder (input_ids + input_ids_2 -> concatenated penultimate hidden states 'encoder_hidden_states' + pooled 'text_embeds') and the SDXL UNet (sample/timestep/encoder_hidden_states/text_embeds/ time_ids -> noise_pred). build_diffusion_pipeline_metadata gains text_encoder_edges to route both conditioning outputs; convert.py emits them for SDXL. time_ids is an external denoiser input the driver supplies; multi-input CFG (onnx-genai) guides both encoder_hidden_states and text_embeds. Verified: mobius convert-comfyui on tiny SDXL produces the dual-edge pipeline + 5-input denoiser + dual-io text encoder. 51 integration tests pass (SDXL dual conditioning-edge routing in metadata + convert).
So the run_comfyui.py driver can branch on SDXL (dual tokenizers + time_ids + dual-conditioning uncond).
ComfyUI LoraLoader nodes are handled by fusing the LoRA into the base model before ONNX export (no runtime LoRA support needed). The translator collects stacked LoraLoaders along the KSampler.model chain (base-first order, with strength_model); checkpoint_export gains a loras=[(path,strength)] arg that runs diffusers load_lora_weights + fuse_lora before export; convert_comfyui_workflow resolves ComfyUI LoRA filenames via a lora_paths map (or absolute paths); the mobius convert-comfyui CLI gains --lora NAME=PATH (repeatable). Validated (scripts/lora_e2e.py, onnx-genai): a fused export matches diffusers load+fuse to 7.9e-6 and differs from the base by 0.63 (LoRA applied). 53 integration tests pass (LoRA collection: stacked order, strengths, none).
Like SDXL, ControlNet fits the declarative pipeline with NO runtime change: the denoiser is a fused ControlNet+UNet taking an extra constant controlnet_cond image input (the ControlNet produces down/mid residuals injected into the UNet). checkpoint_export gains controlnet=<dir|file|__from_unet__>; the translator collects ControlNetApply(control_net_name, strength); convert_comfyui_workflow resolves it via controlnet_paths; mobius convert-comfyui --controlnet NAME=PATH. controlnet_cond is an external denoiser input (like SDXL time_ids) so multi-input CFG shares it across cond/uncond. Validated (scripts/controlnet_e2e.py, onnx-genai): a fused ControlNet+UNet export matches diffusers (unet+controlnet residuals) to 5.8e-6 and differs from the base UNet by 0.45 (ControlNet has effect). 55 integration tests pass.
Compose the SDXL dual-encoder and ControlNet residual paths (removing the earlier 'not supported' guard): _export_sdxl_controlnet fuses ControlNet + SDXL UNet into a denoiser taking sample/timestep/encoder_hidden_states/text_embeds/time_ids/ controlnet_cond (the ControlNet also consumes SDXL's added_cond_kwargs). No new model needed. Validated (scripts/sdxl_controlnet_e2e.py): fused export matches diffusers to 1.1e-6 and differs from the no-controlnet SDXL UNet by 0.16. 55 integration tests pass.
An inpainting UNet takes 9 input channels (latent[4]+mask[1]+masked_latent[4]) and predicts the 4-ch noise. Like SDXL/ControlNet it fits the pipeline with NO runtime change: export_checkpoint auto-detects it (in_channels == out_channels + 5) and _export_sd_inpaint emits a denoiser taking the 4-ch loop-carried sample plus two constant inputs mask/masked_latent, concatenated for the UNet. The latent dim is the UNet out_channels (4), not the 9 input channels. Validated (scripts/inpaint_e2e.py, onnx-genai): a constructed 9-ch UNet (conv_in 4->9) exports and matches the torch UNet to 1.2e-5 while mask/masked_latent change the output by 0.84. 55 integration tests pass.
Export vae_encoder.onnx (image->latent) for inpainting models so a driver can compute the masked-image latent; add exported.inpaint to convert run.json so the renderer branches on it.
…xponential) SchedulerConfig.use_exponential_sigmas (emitted when True); from_diffusers reads it; ComfyUI scheduler='exponential' enables it; convert computes exponential timesteps (fractional -> float32 timestep). 56 integration tests pass.
EmptyLatentImage.batch_size is now recovered into ComfyUIWorkflow.batch_size. The onnx-genai runtime already handles batch>1 (loop-carried state / timestep injection / schedulers are all shape-generic — verified by scripts/batch_check.py). Driver batching (batched prompts/seeds) is a follow-up. 57 integration tests pass.
Co-authored-by: Copilot <223556219+Copilot@users.noreply.github.com> Copilot-Session: bd351e74-2569-444d-a7da-6456f3b5df82
…ckage export_checkpoint now writes a fast-format tokenizer.json (and tokenizer_2.json for SDXL) into the package root, converting the diffusers (possibly slow) CLIP tokenizer to fast. This makes the exported ONNX package self-contained so the native run_comfyui runner can tokenize prompts without the HF cache (which only ships vocab.json + merges.txt). Add checkpoint_export_test.py. Co-authored-by: Copilot <223556219+Copilot@users.noreply.github.com>
The onnx-genai integration now dispatches through write_onnx_genai_config (LLM decoder metadata + diffusion pipeline) rather than the older write_inference_metadata entry point. Co-authored-by: Copilot <223556219+Copilot@users.noreply.github.com> Copilot-Session: bd351e74-2569-444d-a7da-6456f3b5df82
Add build_language_diffusion_pipeline_metadata for masked (discrete) language-diffusion models (LLaDA/Dream): emits an iterative pipeline with a masked_diffusion scheduler (mask_token_id + optional block_length/temperature), a loop-carried logits->input_ids self-edge, and optional unsupervised CFG (guidance_scale). Exported from the package; 5 new tests including schema validation against onnx-genai's committed inference_metadata schema. Co-authored-by: Copilot <223556219+Copilot@users.noreply.github.com> Copilot-Session: bd351e74-2569-444d-a7da-6456f3b5df82
… convert Mobius builds models from scratch (onnx_ir/components), never via torch.onnx. Delete the diffusers->torch.onnx.export shortcut (checkpoint_export.py + test) and make the ComfyUI convert path torch-free: it now only translates the workflow into inference_metadata.yaml + run.json, reconciling the ComfyUI sampler with the checkpoint's diffusers scheduler config (betas). The ONNX component graphs (denoiser/VAE/text encoder) are produced by Mobius's from-scratch diffusers builder (build_diffusers_pipeline) + write_onnx_genai_config. - convert.py: drop export_checkpoint/ExportedCheckpoint; take a SchedulerConfig. - __main__ convert-comfyui: drop --lora/--controlnet/--opset (export-only), add --sdxl. - convert_test: use SchedulerConfig. 62 onnx_genai tests pass. LoRA/ControlNet/inpaint become from-scratch builder capabilities (follow the mobius adding-a-new-model / diffusion-models skills), not export-time fusion. Co-authored-by: Copilot <223556219+Copilot@users.noreply.github.com> Copilot-Session: bd351e74-2569-444d-a7da-6456f3b5df82
Add an optional linear_class factory to _BasicAttention / _CrossAttentionBlock in the from-scratch diffusion UNet, defaulting to plain Linear (backward- compatible). Passing a LoRALinear factory builds the low-rank adapter branch (base + x@A^T@B^T*scale) with an optional shared gate_holder for per-adapter runtime scalar gating (on/off/blend), mirroring the Phi4MM multi-adapter pattern. This is the foundation for runtime LoRA in the diffusion pipeline: adapters baked as lora_A/lora_B params, switched at run time via a gate input (no re-export). New models/unet_lora_test.py (3 graph-build tests, no weights): plain vs LoRA branch, adapter param naming, and runtime gate wiring. Diffusers-builder tests unaffected (21 pass). Co-authored-by: Copilot <223556219+Copilot@users.noreply.github.com> Copilot-Session: bd351e74-2569-444d-a7da-6456f3b5df82
…inputs
Thread the LoRALinear factory from UNet2DConditionModel through the down/mid/up
blocks (all cross-attention projections), driven by a new UNet2DConfig
lora_adapters field ((name, rank, scale) tuples; empty = plain, backward-
compatible). The UNet forward populates a shared _lora_gates mapping, and the
DenoisingTask declares one scalar lora_gate.{name} input per adapter (only for
modules that declare adapters, so DiT is unaffected). At run time a loaded model
switches/blends LoRAs by feeding the gate scalars (1.0/0.0/strength) — no
rebuild; onnx-genai feeds them as denoiser.lora_gate.{name} per-request inputs.
Tests: unet_lora_test.py grows to 5 (full-UNet end-to-end gate-input + adapter-
param assertions, plus no-LoRA has no gate inputs). 141 builder/tasks/unet tests
pass; DiT denoising-task tests green.
Co-authored-by: Copilot <223556219+Copilot@users.noreply.github.com>
Copilot-Session: bd351e74-2569-444d-a7da-6456f3b5df82
Add remap_diffusers_unet_lora + load_unet_lora_safetensors: map a diffusers/
PEFT UNet LoRA state dict (lora.down/up or lora_A/B, under transformer_blocks.N)
onto the from-scratch UNet's baked LoRALinear params (lora_A.{name}.weight /
lora_B.{name}.weight, flattened attn1/attn2), verified against the real built
UNet's param names. Exported from mobius.models.
Tests grow to 8: key remap (both source spellings), cross-check that remapped
keys land on real baked params, and safetensors load w/ shape check. 139
builder/tasks/unet tests pass.
Co-authored-by: Copilot <223556219+Copilot@users.noreply.github.com>
Copilot-Session: bd351e74-2569-444d-a7da-6456f3b5df82
…me LoRA
Wire the LoRA weight-loading into the from-scratch build: unet_loras maps
adapter_name -> lora.safetensors; each is loaded, remapped onto the baked
LoRALinear params, its rank inferred from the file, and injected as
UNet2DConfig.lora_adapters before building the UNet + merged alongside the base
weights. Baked scale = 1.0 (the runtime lora_gate.{name} input supplies the
effective strength). Adds _prepare_unet_loras (unit-tested: rank inference +
merge). 22 builder tests pass.
Runtime LoRA is now code-complete: gated graph (unet.py) + weight load/remap +
build integration; onnx-genai feeds denoiser.lora_gate.{name} per request.
Co-authored-by: Copilot <223556219+Copilot@users.noreply.github.com>
Copilot-Session: bd351e74-2569-444d-a7da-6456f3b5df82
Add an end-to-end test that builds a tiny random-weight LLaDA package via MaskedDiffusionTask and asserts the exported ONNX I/O signature matches the onnx-genai denoiser contract: single input input_ids [B,S] int64, single output logits [B,S,V] f32, and no past/present KV ports. It then verifies build_language_diffusion_pipeline_metadata emits the matching pipeline (denoiser model, logits -> input_ids self-edge, iterative strategy, masked_diffusion scheduler_config with mask_token_id 126336). Co-authored-by: Copilot <223556219+Copilot@users.noreply.github.com> Copilot-Session: bd351e74-2569-444d-a7da-6456f3b5df82
…n names, max_length guard Three fixes found while building a real Stable-Diffusion package (OFA-Sys/small-stable-diffusion-v0) from scratch for onnx-genai: - _diffusers_builder: load component weights from PyTorch .bin (index or single-file) in addition to safetensors — many real diffusers repos ship only .bin. Loader dispatches on extension (safetensors.torch vs torch.load). - models/vae: preprocess_weights now maps legacy diffusers VAE mid-block attention names (query/key/value/proj_attn) to the current to_q/to_k/to_v/to_out.0, squeezing legacy 1x1-conv projections to Linear. Older SD VAE checkpoints (incl. small-stable-diffusion-v0) use the old names. - __main__: guard args.max_length with getattr (the arg is not defined on the build parser; the validation dereferenced a missing attribute and crashed every build --runtime onnx-genai invocation). Co-authored-by: Copilot <223556219+Copilot@users.noreply.github.com> Copilot-Session: bd351e74-2569-444d-a7da-6456f3b5df82
Register the classic SD components so `build_diffusers_pipeline` builds a full text->image package (text_encoder + unet + vae) instead of only the VAE: - Register `UNet2DConditionModel` (denoising) and `CLIPTextModel` (feature-extraction) in `_DIFFUSERS_CLASS_MAP`. - Add `CLIPTextConfig.from_diffusers`, mapping a transformers CLIP text-encoder config onto `ArchitectureConfig` (layer_norm_eps -> rms_norm_eps, derived head_dim, quick_gelu). - Teach the from-scratch UNet to honor `down_block_types`/`up_block_types` so cross-attention lands only where diffusers places it (SD 1.x uses a plain last DownBlock2D and first UpBlock2D). Add the fields to `UNet2DConfig`. - Derive diffusion-pipeline metadata component filenames from the built package keys in `write_onnx_genai_config` (denoiser=unet, text_encoder prompt_only, vae_decoder final_only, latent port `latent_sample`), defaulting CFG to 7.5 when a text encoder is present. Parity vs diffusers fp32 on OFA-Sys/small-stable-diffusion-v0 (built fp16): unet max|Δ|=1.0e-4 (rel 2e-5), text_encoder rel=1.2e-3, vae_decoder 4e-5. Adds an SD 1.x mixed-block-type UNet parity regression test. Co-authored-by: Copilot <223556219+Copilot@users.noreply.github.com> Copilot-Session: bd351e74-2569-444d-a7da-6456f3b5df82
…registration The classic Stable Diffusion integration registered UNet2DConditionModel and CLIPTextModel in _DIFFUSERS_CLASS_MAP but left four _diffusers_builder tests asserting the old (DiT-only) roster, so the suite was red: - test_populates_expected_classes: add UNet2DConditionModel + CLIPTextModel. - test_task_names_are_valid: allow 'feature-extraction' (a real TASK_REGISTRY task, used by the CLIP text encoder). - test_skips_unregistered_class_names: use T5EncoderModel (genuinely unregistered) instead of CLIPTextModel (now registered) to keep exercising the skip/raise path. - test_multiple_components_built: text_encoder (CLIPTextModel) is now a supported component and builds; assert it IS in the result. All 22 _diffusers_builder tests pass; broader diffusion suite green; ruff clean. Co-authored-by: Copilot <223556219+Copilot@users.noreply.github.com> Copilot-Session: bd351e74-2569-444d-a7da-6456f3b5df82
The CLIPTextConfig and UNet2DConfig from_diffusers adapters added for the from-scratch classic Stable Diffusion build had no dedicated test. Adds unit coverage for the non-trivial CLIP mappings (transformers layer_norm_eps -> rms_norm_eps, head_dim = hidden_size/num_attention_heads, num_key_value_heads == num_attention_heads, quick_gelu default) across SD1.5/SD2.x shapes and defaults, plus core UNet field mapping. 4 tests pass; ruff clean. Co-authored-by: Copilot <223556219+Copilot@users.noreply.github.com> Copilot-Session: bd351e74-2569-444d-a7da-6456f3b5df82
Populate Gemma3 GGUF sliding attention config and RoPE scaling metadata, map Gemma3-specific GGUF norm tensors, and keep Gemma3 projections quantized when building with --keep-quantized. Co-authored-by: Copilot <223556219+Copilot@users.noreply.github.com>
![github-code-quality[bot]](https://avatars.githubusercontent.com/u/9919?s=60&v=4){kind=small}
|
|
||
| def _cmd_convert_comfyui(args: argparse.Namespace) -> None: | ||
| """Execute the 'convert-comfyui' subcommand.""" | ||
| import json |
| _SAMPLER_NODES = ("KSampler", "KSamplerAdvanced") | ||
| _VAE_DECODE_NODES = ("VAEDecode", "VAEDecodeTiled") | ||
| _TEXT_ENCODE_NODES = ("CLIPTextEncode",) | ||
| _CHECKPOINT_NODES = ("CheckpointLoaderSimple", "CheckpointLoader", "UNETLoader") |
Contributor
There was a problem hiding this comment.
Pull request overview
This PR expands Mobius’ model/task/runtime integration surface to better support GGUF-imported Gemma3 (scaled/local RoPE, hybrid attention schedules, quantized builds), Stable Diffusion–style diffusers pipelines (including runtime-gated UNet LoRA), and onnx-genai inference metadata generation/conversion utilities.
Changes:
- Fix and extend GGUF →
ArchitectureConfig/ tensor-name mapping for Gemma3 (sliding window, local RoPE base, rope_scaling defaults, norm/QK-norm tensors, quantized component wiring). - Add diffusion-oriented task/model capabilities: runtime-gated LoRA inputs for denoisers, diffusers pipeline build enhancements, and from-scratch UNet/VAE legacy weight remaps + parity tests.
- Introduce onnx-genai
inference_metadata.yamlgeneration for diffusion pipelines, ComfyUI workflow conversion, and add a new MLX EP capability profile.
Reviewed changes
Copilot reviewed 39 out of 39 changed files in this pull request and generated 4 comments.
Show a summary per file
| File | Description |
|---|---|
| src/mobius/tasks/_masked_diffusion.py | New task for masked-diffusion (bidirectional, no-KV) language models. |
| src/mobius/tasks/_denoising.py | Adds optional runtime LoRA gate inputs for denoising tasks. |
| src/mobius/tasks/init.py | Registers the new masked-diffusion task. |
| src/mobius/models/vae.py | Adds legacy diffusers VAE attention weight-name remapping + 1x1-conv squeeze. |
| src/mobius/models/unet.py | Adds runtime LoRA plumbing, block-type awareness, and projection-mode handling for UNet. |
| src/mobius/models/unet_parity_test.py | New numerical parity tests vs diffusers for UNet and attention blocks. |
| src/mobius/models/unet_lora_test.py | New graph-build tests for LoRA branches + runtime gate wiring + key remapping helpers. |
| src/mobius/models/llada.py | New bidirectional masked-diffusion language model implementation (LLaDA). |
| src/mobius/models/llada_test.py | New parity + bidirectionality + masked-diffusion metadata contract tests. |
| src/mobius/models/gemma3_text.py | Quantized component selection + embedding scaling + LM head re-tie fixes. |
| src/mobius/models/gemma3_text_test.py | New tests for quantized component usage + retied LM head. |
| src/mobius/models/init.py | Exports LLaDA + UNet LoRA helper functions. |
| src/mobius/integrations/ort_genai/genai_config_test.py | Adds MLX EP test coverage for shared KV behavior. |
| src/mobius/integrations/onnx_genai/inference_metadata.py | Major refactor: pipeline-style metadata for diffusion + scheduler config loading/parsing. |
| src/mobius/integrations/onnx_genai/inference_metadata_test.py | Rewrites tests around new diffusion + language-diffusion metadata emitters. |
| src/mobius/integrations/onnx_genai/decoder_metadata.py | New module for decoder-only LLM inference metadata generation/writing. |
| src/mobius/integrations/onnx_genai/decoder_metadata_test.py | New tests for decoder metadata and schema validation. |
| src/mobius/integrations/onnx_genai/convert.py | New ComfyUI → onnx-genai conversion orchestration + timestep best-effort. |
| src/mobius/integrations/onnx_genai/convert_test.py | New tests for conversion reconciliation behavior. |
| src/mobius/integrations/onnx_genai/comfyui.py | New ComfyUI workflow parser/translator producing pipeline metadata + run params. |
| src/mobius/integrations/onnx_genai/comfyui_test.py | New tests for ComfyUI parsing/translation and schema validation. |
| src/mobius/integrations/onnx_genai/auto_export.py | New dispatcher to write onnx-genai metadata based on package type. |
| src/mobius/integrations/onnx_genai/auto_export_test.py | Tests for the auto-export dispatcher. |
| src/mobius/integrations/onnx_genai/init.py | Public API surface update for new onnx-genai helpers. |
| src/mobius/integrations/gguf/_tensor_mapping.py | Adds Gemma3-specific tensor-name mapping extras (norms, q/k norms). |
| src/mobius/integrations/gguf/_tensor_mapping_test.py | Adds Gemma3 mapping regression tests. |
| src/mobius/integrations/gguf/_config_mapping.py | Extends GGUF config mapping for sliding window, local RoPE base, rope_scaling defaults, Gemma3 activation/QK norm. |
| src/mobius/integrations/gguf/_config_mapping_test.py | New synthetic GGUF metadata tests for rope_scaling/local RoPE/layer schedule. |
| src/mobius/_registry.py | Registers model_type="llada" with the masked-diffusion task. |
| src/mobius/_model_package.py | Updates save() doc reference to new onnx-genai export entrypoint. |
| src/mobius/_exporter_test.py | Adds MLX EP capability test coverage. |
| src/mobius/_execution_providers.py | Adds MLX EP capability registration (GQA dtypes, shared KV, no QKV packing). |
| src/mobius/_diffusers_configs.py | Adds CLIP text config adapter + expands UNet config (block types, projection mode, LoRA adapters). |
| src/mobius/_diffusers_configs_test.py | New tests for CLIPTextConfig and UNet2DConfig adapters. |
| src/mobius/_diffusers_builder.py | Adds classic SD component support, .bin weight loading fallback, and UNet LoRA baking/merge. |
| src/mobius/_diffusers_builder_test.py | Updates coverage for new SD components and LoRA preparation helper. |
| src/mobius/_configs/_base.py | Adds more transformer-config field aliases (activation_type, d_model, n_heads, etc.). |
| src/mobius/main.py | Updates onnx-genai export path + adds convert-comfyui CLI subcommand. |
| docs/execution_providers.md | Documents the new MLX execution provider. |
Suppressed comments (2)
src/mobius/models/unet.py:391
- Same naming ambiguity as the down block:
attention_head_dimis used asnum_heads. A brief comment would make the intended semantics explicit.
src/mobius/models/unet.py:480 attention_head_dimis being treated as the number of heads in this UNet mid block. Consider clarifying with an inline comment to prevent confusion withdim_headterminology used elsewhere (e.g. Transformer2DModel).
Comment on lines
+112
to
+119
| cfg = config if config is not None else getattr(pkg, "config", None) | ||
| if cfg is None: | ||
| raise ValueError( | ||
| "onnx-genai decoder metadata requires a model config (pass config=... " | ||
| "or a package carrying `.config`)" | ||
| ) | ||
| path = write_decoder_metadata(output_dir, config=cfg, kv_native_dtype=kv_native_dtype) | ||
| return {"inference_metadata": path} |
Comment on lines
+77
to
+81
| Ancestral samplers (which inject fresh noise every step) have no | ||
| deterministic onnx-genai equivalent and are rejected, as are scheduler | ||
| classes onnx-genai does not implement, so a Mobius-built package never | ||
| silently runs the wrong denoise dynamics. | ||
| """ |
Comment on lines
+99
to
+102
| raise ValueError( | ||
| f"unsupported diffusers scheduler {raw_name!r} for onnx-genai; " | ||
| f"supported kinds: ddim (DDIMScheduler), euler (EulerDiscreteScheduler)" | ||
| ) |
| ) | ||
| if cross_attention_dim: | ||
| num_heads = max(1, out_channels // attention_head_dim) | ||
| num_heads = attention_head_dim |
Repack independent GGUF tensors with a bounded worker pool, reduce Q4 nibble-reorder temporaries, and show packed/written byte progress during long conversions. Co-authored-by: Copilot <223556219+Copilot@users.noreply.github.com>
Member
Author
|
Performance update added in commit What changed:
Profiling/measurement before optimizing:
After changes:
Gemma-3 27B rebuild after optimization:
Validation after perf changes:
|
Wire --max-shard-size through build-gguf and save sharded safetensors external data with bounded parallel shard writers while preserving ONNX external-data loadability. Co-authored-by: Copilot <223556219+Copilot@users.noreply.github.com>
Member
Author
|
Follow-up on sharding/write-side question, added in commit Findings from reading the code:
What changed:
Important route decision for the 27B conversion:
Updated validation:
|
This reverts commit 389bb5d.
Member
Author
|
Update on ONNX external-data sharding and the safetensors detour: I verified the installed
So the honest current status is: ONNX external-data sharding may exist in unreleased onnx-ir work, but it is not present in the released 0.2.1 API available to this mobius environment. Per the updated direction, I did not pursue safetensors as the solution and pushed What remains in this PR:
Validation after the revert:
|
Wire max_shard_size_bytes through the ONNX external-data save path and expose --max-shard-size for GGUF builds. Add a runtime error for onnx-ir releases that predate onnx/ir-py PR #427, plus regression coverage for passing the option and writing sharded ONNX data. Co-authored-by: Copilot <223556219+Copilot@users.noreply.github.com>
Member
Author
|
Update after installing onnx-ir from git main:
Validation:
Operational note: onnx-ir's sharded writer raises |
Member
Author
|
Final artifact/loadability update: After the latest correction, I stopped treating unreleased ONNX external-data sharding as a prerequisite for the artifact. The model now left on disk is plain ONNX external data with a single
Verification:
Full Sharding note for reviewers: onnx/ir-py PR #427 is merged but not in PyPI's onnx-ir 0.2.1. It adds |
This was referenced Aug 6, 2026
Treat GGUF Q4_0 and Q8_0 as requiring explicit zero-point initializers so MatMulNBits and GatherBlockQuantized do not rely on execution-provider-specific defaults. Add a GatherBlockQuantized dequantization regression that checks the zero-point value, not only node wiring. Co-authored-by: Copilot <223556219+Copilot@users.noreply.github.com>
Member
Author
|
Added the zero-point fix to this PR in commit What changed:
Regression coverage:
Quantization types checked:
Validation:
Upstream ORT issue:
|
Member
Author
|
Closing this PR rather than merging/rebasing it: it was accidentally based on another in-flight feature branch and includes unrelated commits. I rebuilt the urgent zero-point correctness fix from current origin/main in #459. The remaining Gemma3 config and repack performance work should be reapplied semantically as separate clean-main PRs. |
justinchuby
added a commit
that referenced
this pull request
Aug 6, 2026
## Summary This is the clean-main rebuild of the urgent correctness fix from #458. It contains only the GGUF zero-point fix, based on current `origin/main`. GGUF Q4_0 and Q8_0 are symmetric formats, but their dequantization formulas still require non-zero zero points: - Q4_0: `(q - 8) * scale` - Q8_0: `(q - 128) * scale` Mobius previously marked those GGUF imports as symmetric, so `QuantizedLinear` and `QuantizedEmbedding` omitted the optional `zero_points` input for `MatMulNBits` and `GatherBlockQuantized`. That is not portable: `GatherBlockQuantized` defaults diverge between ORT CPU and ORT CUDA when `zero_points` is omitted, corrupting CUDA embeddings before the first decoder layer runs. This PR makes GGUF Q4_0/Q8_0 emit explicit zero-point initializers, so both `GatherBlockQuantized` and `MatMulNBits` receive the intended values instead of relying on EP defaults. ## Regression test Added a single-node `GatherBlockQuantized` runtime regression: - qweight nibbles are all `10` - scales are known constants - explicit zero point is `8` - output is asserted against hand-computed `(q - 8) * scale` - rerunning with zero point `0` is asserted not to match I also temporarily flipped the test's zero point to `0`; it failed with a 100% output mismatch, so the test checks the value and not just the presence of the input. Synthetic GGUF tests now also assert: - Q4_0 `MatMulNBits` nodes have the explicit fourth input - Q4_0 `GatherBlockQuantized` has the explicit fourth input - Q4_0 embedding zero points are packed as `0x88` ## Quantization types checked Direct GGUF repackers already produce explicit zero-points for supported direct formats: - Q4_0: zp=8 - Q4_1: per-block affine zp - Q4_K: requantized per-block affine zp - Q8_0: zp=128 - Q1_0: zp=1; Tencent Q1_0 uses its custom zp path Q5_0/Q5_1/Q5_K/Q6_K are not direct `MatMulNBits` repack targets here; when they appear in mixed presets they go through dequantize+requantize/native-block fallback paths, which produce explicit zero-points where needed. ## Validation - `python -m pytest src\mobius\integrations\gguf\_builder_test.py -q` -> `33 passed` - `python -m pytest src\mobius\integrations\gguf\ src\mobius\_configs\ src\mobius\_model_package_test.py` -> `304 passed` - `python -m ruff check src\mobius\integrations\gguf\_builder.py src\mobius\integrations\gguf\_builder_test.py` -> passed - `python -m ruff format --check src\mobius\integrations\gguf\_builder.py src\mobius\integrations\gguf\_builder_test.py` -> passed Fresh conversion validated with `C:\Users\justinchu\dev\models-gguf\qwen2.5-0.5b-instruct-q4_0.gguf`: ```powershell python -m mobius build-gguf C:\Users\justinchu\dev\models-gguf\qwen2.5-0.5b-instruct-q4_0.gguf --output C:\Users\justinchu\dev\models\qwen2.5-0.5b-q4_0-mobius --keep-quantized --dtype f16 --ep cuda ``` Converted graph verification: - `GatherBlockQuantized` input count: 4 - all `MatMulNBits` input counts: 4 - zero-point initializers: 170 Generation checks on the freshly converted model, no post-hoc graph patching: - ORT CPU via onnx-genai CLI: `The capital of France is Paris. It is the largest city in` - ORT CUDA via onnx-genai CLI (`ONNX_GENAI_ORT_LIB_DIR` pointed at the installed onnxruntime-gpu package, CUDA/cuDNN DLL dirs on PATH): `The capital of France is Paris. It is the largest city in` - native CUDA via onnx-genai CLI (`--features native-cuda`, CUDA/cuDNN DLL dirs on PATH): `The capital of France is Paris. It is the largest city in` ## Upstream ORT issue Reported the provider default divergence here: microsoft/onnxruntime#31692 ## Relationship to #458 #458 was accidentally based on another in-flight feature branch and includes unrelated commits. This PR is the clean-main replacement for the urgent zero-point correctness fix only. --------- Signed-off-by: Justin Chu <justinchuby@users.noreply.github.com> Co-authored-by: Copilot <223556219+Copilot@users.noreply.github.com> Co-authored-by: copilot-swe-agent[bot] <198982749+Copilot@users.noreply.github.com> Co-authored-by: justinchuby <11205048+justinchuby@users.noreply.github.com>
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
Fix GGUF -> ArchitectureConfig gaps that surface as deep model-construction errors instead of diagnosable GGUF metadata issues. These are general GGUF mapping bugs, not one-checkpoint quirks: any scaled-RoPE GGUF can set
rope_typetolinear/yarn/longropewhile leavingconfig.rope_scalingunset, and any hybrid/sliding-window architecture needs the layer schedule and local RoPE fields populated before graph construction.Changes:
rope_scaling(includingfactor) when GGUF declaresrope.scaling.typeaslinear,yarn, orlongrope; warn and use identity factor1.0if the factor key is absent instead of failing later.rope.freq_base_swa, or default to10000.0for sliding-window GGUFs (matching HF Gemma3 defaults).attention.sliding_windowand derive Gemma3layer_typesfrom HF's architectural defaultsliding_window_pattern=6(five sliding-window layers, then one full-attention layer).--keep-quantizedproducesMatMulNBitsinstead of expanding projections to f16.Errors this prevents
These failures previously occurred far from GGUF parsing and did not name the missing metadata/config field:
TypeError: 'NoneType' object is not subscriptablefromLinearRope.__init__whenrope_typewas scaled butconfig.rope_scalingwasNone.TypeError: unsupported operand type(s) for **: 'NoneType' and 'float'from_get_default_inv_freqwhen sliding-window local RoPE theta was missing.TypeError: 'NoneType' object is not iterablefromgemma3_text.pywhenconfig.layer_typeswas missing.ValueError: Component 'model' has 124 initializer(s) without weights...after graph construction when Gemma3 GGUF norm tensors were skipped/mis-mapped.--keep-quantizedbuilt a ~53 GiB f16 model withMatMulNBits=0; after this change the same build is ~16 GiB withMatMulNBits=434.I considered raising explicit config errors for absent
rope.scaling.factor, but kept the warning + identity fallback because1.0is a safe no-scaling factor and lets partially specified GGUF metadata remain buildable. The missing local RoPE and layer schedule are now populated before the model code can fail with genericTypeErrors.Regression tests
Added synthetic metadata/unit tests; no large GGUF is required. I also verified they fail without their corresponding implementation changes:
_config_mapping.pyand runningpython -m pytest src\mobius\integrations\gguf\_config_mapping_test.py -qfails all 5 new config tests (rope_scaling, local RoPE, sliding window/layer types)._tensor_mapping.pyand runningpython -m pytest src\mobius\integrations\gguf\_tensor_mapping_test.py -qfails the Gemma3 norm mapping regression.gemma3_text.pyand runningpython -m pytest src\mobius\models\gemma3_text_test.py -qfails the quantized-component and retied-head regressions.Validation
python -m ruff check src\mobius\integrations\gguf\_config_mapping.py src\mobius\integrations\gguf\_config_mapping_test.py src\mobius\integrations\gguf\_tensor_mapping.py src\mobius\integrations\gguf\_tensor_mapping_test.py src\mobius\models\gemma3_text.py src\mobius\models\gemma3_text_test.pypython -m ruff format --check src\mobius\integrations\gguf\_config_mapping.py src\mobius\integrations\gguf\_config_mapping_test.py src\mobius\integrations\gguf\_tensor_mapping.py src\mobius\integrations\gguf\_tensor_mapping_test.py src\mobius\models\gemma3_text.py src\mobius\models\gemma3_text_test.pypython -m pytest src\mobius\integrations\gguf\ src\mobius\_configs\ src\mobius\models\gemma3_text_test.py->182 passedpython -m mobius build-gguf C:\Users\justinchu\dev\models-gguf\gemma-3-27b-it-Q4_0.gguf --output C:\Users\justinchu\dev\models\gemma-3-27b-onnx --keep-quantized --dtype f16 --ep cudaC:\Users\justinchu\dev\models\gemma-3-27b-onnx\model.onnx16.079 GiBonnx.checker.check_model(...): OKMatMulNBits: 434GatherBlockQuantized: 0 (token_embd.weightis Q6_K in this GGUF, so it does not match the graph's Q4_0 projection quantization config and remains f16)Why this matters
This came up while preparing a 27B int4 model (~15-16 GiB) to test virtual-memory management, weight offloading, weight streaming, and prefetch in
onnx-genaion an 8 GB RTX 4060. Prior measurements used qwen2.5-0.5b, whose weights fit comfortably on the card and therefore cannot exercise PCIe saturation or memory-manager pressure.