Inference-Driven Model Compiler

Export 🤗 Transformers models to ONNX by observing a real inference pass instead of relying on hand-written, per-architecture ONNX configurations.

Standard Optimum ONNX export requires a model-specific OnnxConfig that declares every input/output and which tensor dimensions are dynamic. This project removes that requirement: it runs the model on your actual inputs, traces the tensor shapes that flow through it, determines which dimensions are dynamic empirically, and exports the result — all behind the familiar from_pretrained(...) interface.

It is built entirely on top of an unmodified optimum / optimum-onnx installation.

---

Installation

pip install torch transformers onnx onnxruntime
pip install "optimum @ git+https://github.com/huggingface/optimum"
# the ONNX exporter/runtime now lives in the separate optimum-onnx package
pip install "optimum-onnx[onnxruntime] @ git+https://github.com/huggingface/optimum-onnx"

git clone https://github.com/naomili0924/inference_driven_model_compiler.git
# Put the *repo directory itself* on PYTHONPATH. This activates the shadow
# `optimum` package, which (a) makes `optimum-cli` use the inference-driven
# exporter and (b) auto-registers the new export flags onto
# `optimum-cli export onnx` (see optimum/commands/register/register_idmc.py).
# Add the parent dir too if you also want `import inference_driven_model_compiler`
# or the `idmc` CLI.
export PYTHONPATH=/path/to/inference_driven_model_compiler:$PYTHONPATH

With this repo off PYTHONPATH, optimum-cli behaves exactly as stock — the integration is inert unless the shadow optimum is active.

---

CLI export

When this repo is on PYTHONPATH (see Installation), the standard optimum-cli export onnx command gains four extra flags — no separate tool or launcher needed. They are registered automatically via optimum/commands/register/register_idmc.py, which optimum-cli auto-discovers:

--fixed_inputs — value-exact inputs: by default the tracer records only the shape of each traced input and regenerates a random tensor of that shape at export/validation time (which also lets it discover dynamic axes). That is fine when only the shape matters, but breaks for inputs whose value drives graph construction — sizes, counts, indices, or anything read via .item() / .tolist() / a reshape target — or inputs coupled to another (e.g. a vision encoder's image_grid_thw, whose product must equal pixel_values' patch count). List those names here to replay their real traced values verbatim:

--fixed_inputs='["image_grid_thw", "cache_position"]'

(For image-text-to-text models the vision encoder's pixel_values / image_grid_thw are pinned automatically; use --fixed_inputs for any additional value-exact inputs.)

JSON-or-file: --inference_kwargs, --module_fixed_axis_fields and --fixed_inputs accept either inline JSON or a path to a .json file — handy when the inputs are large (e.g. a multimodal inference_kwargs with a big pixel_values array). Three accepted forms, resolved in this order:

--inference_kwargs=@inputs.json     # '@' prefix: read from file (curl-style)
--inference_kwargs=inputs.json      # bare path to an existing file: read from file
--inference_kwargs='{"input_ids": []}'   # inline JSON string

The @ prefix only forces file interpretation; a bare path works too as long as the file exists at parse time.

Encoder model

optimum-cli export onnx \
    --model sentence-transformers/paraphrase-MiniLM-L12-v2 \
    /dev/shm/paraphrase-MiniLM \
    --export_by_inference=true \
    --module_fixed_axis_fields='{"transformer": ["hidden_size","intermediate_size","type_vocab_size","vocab_size"]}'

Decoder model (with KV cache)

optimum-cli export onnx \
    --model Qwen/Qwen3-4B-Thinking-2507 \
    /dev/shm/qwen3-4b-thinking-onnx \
    --task text-generation-with-past \
    --export_by_inference=true \
    --dtype fp16 --device cpu

The inference-driven tracer builds its dummy inputs on CPU, so export decoder models with --device cpu (a --device cuda model would mismatch the CPU-resident traced inputs).

--module_fixed_axis_fields is optional for decoder models — the dynamic-axis inference step figures out num_heads, head_dim, etc. automatically.

Tip for large models: if your disk is limited, export to /dev/shm (a RAM-backed tmpfs typically >= 80 GB on GPU instances) and copy the result elsewhere afterwards.

---

Python API

Encoder model (BERT feature extraction)

from transformers import AutoTokenizer
from inference_driven_model_compiler.optimum.onnxruntime import (
    OnTheFlyORTModelForFeatureExtraction,
)

ckpt = "bert-base-uncased"
tokenizer = AutoTokenizer.from_pretrained(ckpt)
encoded = tokenizer("ONNX Runtime accelerates inference.", return_tensors="pt")

model = OnTheFlyORTModelForFeatureExtraction.from_pretrained(
    ckpt,
    inference_kwargs=dict(encoded),
    export_by_inference=True,
    export=True,
    module_fixed_axis_fields={"transformer": ["hidden_size", "num_attention_heads"]},
)

out = model(**encoded)
print(out.last_hidden_state.shape)      # (1, seq_len, 768)

Decoder model (GPT-2 text generation, with KV cache)

from transformers import GPT2Tokenizer
from inference_driven_model_compiler.optimum.onnxruntime import OnTheFlyORTModelForCausalLM

ckpt = "gpt2"
tokenizer = GPT2Tokenizer.from_pretrained(ckpt)
encoded = tokenizer("Replace me by any text you'd like.", return_tensors="pt")

model = OnTheFlyORTModelForCausalLM.from_pretrained(
    ckpt,
    inference_kwargs=dict(encoded),
    export_by_inference=True,
    export=True,
    module_fixed_axis_fields={"transformer": ["n_ctx", "n_embd"]},
)

output_ids = model.generate(**encoded)
print(tokenizer.decode(output_ids[0]))

Diffusion pipeline (text-to-video)

Pass export_by_inference=True together with inference_kwargs (the same kwargs you would pass to the pipeline __call__). The pipeline runs once in PyTorch to capture real tensor shapes for every submodule, then exports each one to ONNX automatically — no hand-written OnnxConfig required.

import torch
from inference_driven_model_compiler.optimum.onnxruntime import OnTheFlyORTDiffusionPipeline

inf_kwargs = {
    "prompt": "A cat walks on the grass, realistic",
    "negative_prompt": "low quality, blurred",
    "height": 240,
    "width": 416,
    "num_frames": 21,
    "guidance_scale": 5.0,
}

pipe = OnTheFlyORTDiffusionPipeline.from_pretrained(
    "Wan-AI/Wan2.1-T2V-1.3B-Diffusers",
    provider="CUDAExecutionProvider",
    torch_dtype=torch.float16,
    export_by_inference=True,
    inference_kwargs=inf_kwargs,
    module_fixed_axis_fields={
        "text_encoder": ["d_model", "vocab_size"],
        "transformer":  ["in_channels", "text_dim"],
        "vae_decoder":  ["base_dim", "z_dim"],
    },
)

output = pipe(**inf_kwargs).frames[0]

This exports three ONNX files to a temporary directory and immediately loads them into ORT sessions — all in one from_pretrained call:

Diffusion pipeline (text-to-image)

The same export_by_inference=True path handles UNet-based text-to-image pipelines such as SDXL. The compiler traces the two CLIP text encoders, the UNet (including the pooled added_cond_kwargs micro-conditioning), and the VAE decoder, then runs the whole pipeline through ONNX Runtime.

import torch
from inference_driven_model_compiler.optimum.onnxruntime import OnTheFlyORTDiffusionPipeline

inf_kwargs = {
    "prompt": "A cinematic photo of a red panda astronaut on the moon",
    "num_inference_steps": 1,     # SDXL-Turbo is a one-step distilled model
    "guidance_scale": 0.0,        # no classifier-free guidance
}

pipe = OnTheFlyORTDiffusionPipeline.from_pretrained(
    "stabilityai/sdxl-turbo",
    provider="CUDAExecutionProvider",
    torch_dtype=torch.float16,
    export_by_inference=True,
    inference_kwargs=inf_kwargs,
)

image = pipe(**inf_kwargs).images[0]
image.save("sdxl_turbo_output.png")

This exports four ONNX submodules and immediately loads them into ORT sessions:

The SDXL VAE is numerically unstable in fp16 (the diffusers pipeline upcasts it to fp32 via force_upcast), so the VAE decoder is exported in fp32 even when the rest of the pipeline is fp16. The fp16 latents are auto-cast at inference.

Inference-driven diffusion export traces and exports each submodule on the pipeline's device. Use a CUDA provider for large UNet-based models — exporting an fp16 UNet/VAE on CPU is extremely slow.

Diffusion pipeline (image editing)

Image-editing models — e.g. InstructPix2Pix (timbrooks/instruct-pix2pix) — take an input image plus a text instruction and produce an edited image. OnTheFlyORTImageEditPipeline handles the two things that make them different from text-to-image:

• the input image is encoded to latents through the VAE encoder (vae.encode(image).latent_dist), which text-to-image never touches, and
• the UNet consumes the noisy latents concatenated with the encoded image latents (InstructPix2Pix's UNet has in_channels == 8).

Pass an image (and a prompt) in inference_kwargs so the single tracing pass exercises the VAE encoder and the wide UNet — then every submodule is exported to ONNX at the traced resolution.

import torch
from PIL import Image
from inference_driven_model_compiler.optimum.onnxruntime import OnTheFlyORTImageEditPipeline

inf_kwargs = {
    "prompt": "turn him into a cyborg",
    "image": Image.open("input.png").convert("RGB"),  # any RGB image
    "num_inference_steps": 10,
    "image_guidance_scale": 1.5,
    "guidance_scale": 7.5,
}

pipe = OnTheFlyORTImageEditPipeline.from_pretrained(
    "timbrooks/instruct-pix2pix",
    provider="CUDAExecutionProvider",
    torch_dtype=torch.float32,
    export_by_inference=True,
    inference_kwargs=inf_kwargs,
)

edited = pipe(**inf_kwargs).images[0]
edited.save("edited.png")

This exports four ONNX submodules (note the vae_encoder, exported as the full encode graph — encoder + quant_conv → latent_parameters — so the input image becomes a proper latent distribution):

The concrete diffusers pipeline is resolved from the checkpoint's _class_name and mixed into OnTheFlyORTImageEditPipeline on the fly, so other image-conditioned pipelines (img2img, SDXL InstructPix2Pix, upscaling, …) export and run through the same class without a model-specific subclass. For pure text-to-image, use OnTheFlyORTDiffusionPipeline instead.

InstructPix2Pix uses the SD-1.5 VAE, which is fp16-fragile on some images. The example above exports in fp32 (the model is small, ~1 GB); switch to torch_dtype=torch.float16 for a faster/smaller export once you've confirmed the output quality on your inputs.

Saving & uploading the exported ONNX to the Hub

export_by_inference=True writes the ONNX graphs to ephemeral RAM (/dev/shm) and re-exports on every call. To keep them — and reuse them without re-exporting — save the pipeline to disk and/or push it to the Hugging Face Hub:

# 1. Export once...
pipe = OnTheFlyORTImageEditPipeline.from_pretrained(
    "timbrooks/instruct-pix2pix",
    provider="CUDAExecutionProvider", torch_dtype=torch.float32,
    export_by_inference=True, inference_kwargs=inf_kwargs,
)

# 2. ...then persist locally and/or upload to your Hub repo.
pipe.save_pretrained(
    "instruct-pix2pix-onnx",          # local copy
    push_to_hub=True,
    repo_id="your-username/instruct-pix2pix-onnx",
    token="hf_…",                     # write-scoped token (or use the cached login)
    private=True,
)

save_pretrained writes every submodule's model.onnx (+ external .onnx_data for graphs over 2 GB), the per-submodule config.json, the scheduler/ / tokenizer/ / feature_extractor/, the model_index.json, and the io_binding/ output-shape files — everything needed to reload without tracing.

Loading pre-exported ONNX weights

# From a local directory or a Hub repo — no PyTorch, no re-export.
pipe = OnTheFlyORTImageEditPipeline.from_pretrained(
    "your-username/instruct-pix2pix-onnx",   # or a local path
    export=False,
    provider="CUDAExecutionProvider",
)
edited = pipe(**inf_kwargs).images[0]

# Generic text-to-image works the same way:
pipe = OnTheFlyORTDiffusionPipeline.from_pretrained(
    "optimum/stable-diffusion-v1-5",   # Hub repo with pre-exported ONNX weights
    export=False,
)

Common arguments

---

Available classes

All live in inference_driven_model_compiler.optimum.onnxruntime and share the same from_pretrained(...) interface.

Naming convention: every user-facing entry point carries the OnTheFlyORT prefix — it marks the inference-driven (on-the-fly export) path and keeps these classes from colliding with stock optimum's ORT* names. Internal ONNX submodule wrappers (ORTUnet, ORTVaeEncoder, ORTVae, …) keep the bare ORT prefix. The pre-rename names (ORTDiffusionPipeline, ORTImageEditPipeline, ORTModelForImageTextToText, ORTChatterboxPipeline) remain importable as deprecated aliases, so existing code keeps working.

Transformer models

Diffusion pipelines

OnTheFlyORTDiffusionPipeline requires no model-specific subclass. When called as the base class it reads _class_name from the model's model_index.json and creates an ORT<ClassName> wrapper on the fly via _make_ort_pipeline_class. Every diffusers pipeline — including ones not yet written — is handled automatically.

Supported text-to-video pipeline names (as of diffusers 0.38): AnimateDiffPipeline, AnimateDiffSDXLPipeline, CogVideoXPipeline, HunyuanVideo15Pipeline, HunyuanVideoPipeline, LTXPipeline, LTX2Pipeline, LattePipeline, MochiPipeline, SanaVideoPipeline, TextToVideoSDPipeline, WanPipeline, WanAnimatePipeline.

---

Verified models

Transformer models

Diffusion pipelines

---

Limitations

• Encoder-decoder models (T5, BART, Whisper) are exported encoder-only for feature-extraction; full encoder-decoder generation is not yet wired up.
• CLIP exports the vision encoder (the full CLIP forward needs both text and image inputs and returns embeddings rather than last_hidden_state).
• The exported ONNX is written to a temporary directory; call model.save_pretrained(...) to persist it.
• Diffusion pipeline export runs one full inference pass before exporting, which requires enough GPU/CPU memory to hold the full PyTorch pipeline during tracing.
• VAE encoder export is included in the export spec but the WAN pipeline does not use it during text-to-video inference; it is exported as a no-op placeholder when the submodule exists on the VAE.

---

How it works

Transformer models

from_pretrained(export_by_inference=True)
        │
        ▼
1. Load the PyTorch model (TasksManager)
        │
        ▼
2. trace_model_shapes()  ── run N inference passes with varied input shapes
        │                    • encoder-only      → single forward
        │                    • encoder-decoder   → encoder submodule only
        │                    • decoder-only      → prefill + decode (KV cache)
        ▼
3. _compute_dynamic_axes()  ── a dim is "dynamic" iff its size changed across runs
        │                       (dim-0/batch always dynamic; hidden_size, vocab,
        ▼                        num_heads, head_dim, image H/W … stay static)
4. DummyOnnxConfig  ── a generic OnnxConfig built from the traced shapes + axes
        │
        ▼
5. export_models()  ── standard Optimum ONNX export (disable_dynamic_axes_fix=True)
        │
        ▼
6. ORTModel._from_pretrained()  ── load the ONNX model into an ORT session

Diffusion pipelines

OnTheFlyORTDiffusionPipeline.from_pretrained(export_by_inference=True, inference_kwargs={...})
        │
        ▼
1. Load the PyTorch diffusion pipeline (diffusers)
        │
        ▼
2. Register forward pre-hooks on each submodule
   (text_encoder, transformer/unet, vae.post_quant_conv)
        │
        ▼
3. Run ONE full pipeline inference pass with the provided inference_kwargs
   — all submodule inputs are captured live as real tensors
        │
        ▼
4. For each submodule:
   • Move captured tensors to CPU
   • Build DummyOnnxConfig from the observed shapes + dynamic axes
   • Export to ONNX (constant folding disabled for text encoders to
     avoid 89 GB inflation from precomputed attention bias)
        │
        ▼
5. VAE decoder special case: export post_quant_conv + decoder as a single
   _VaeFullDecodeWrapper (WAN VAE decodes per-frame with caching in PyTorch;
   ONNX needs one call over the full latent video)
        │
        ▼
6. OnTheFlyORTDiffusionPipeline loaded with each ONNX submodule in an ORT session

Dynamic-axis inference

Rather than guessing from config fields, the compiler runs the model several times (default n_trials=3) with randomly varied batch sizes and sequence lengths (and, for decoder models, a varied decode query length). A dimension is marked dynamic only if its value actually changes between runs; everything else — hidden_size, num_attention_heads, head_dim, vocab_size, image height/width, ViT patch count, etc. — is correctly kept static.

For each tensor seen across the trial runs:

• Dimension 0 → always dynamic (batch).
• Any other dimension → dynamic iff its size differed between at least two trials; otherwise static.
• For decoder KV-cache tensors past_key_values.{i}.key/value, the past-sequence dimension (axis 2) is dynamic while num_heads (axis 1) and head_dim (axis 3) stay static.

---

Why

---

Project layout

inference_driven_model_compiler/
├── cli.py                        # idmc CLI — wraps optimum-cli with extra flags
├── optimum/
│   ├── exporters/onnx/
│   │   ├── utils.py              # trace_model_shapes(), dynamic-axis inference
│   │   ├── model_configs.py      # DummyOnnxConfig — generic shape-driven OnnxConfig
│   │   ├── input_generators.py   # DummyTupleInputGenerator — dtype-aware dummies
│   │   └── __init__.py           # main_export wrapper (renamed params)
│   └── onnxruntime/
│       ├── modeling.py           # _OnTheFlyORTMixin + 5 encoder model classes
│       ├── modeling_decoder.py   # OnTheFlyORTModelForCausalLM
│       ├── modeling_diffusion.py # OnTheFlyORTDiffusionPipeline + submodule wrappers
│       └── utils.py              # load_shapes_as_torch_size and helpers
└── on_the_fly_pipeline_tests/    # per-model tests + dynamic-axis + diffusion suites