README.md

metadata

license: unlicense
title: Sdxl onnx
sdk: docker
emoji: 📉
colorFrom: green
colorTo: purple
pinned: true

sd-onnx-cpu-example

This repository contains scripts to:

• Export a Stable Diffusion pipeline (text encoder, UNet, VAE decoder) to ONNX
• Quantize the ONNX models (dynamic quantization)
• Serve image generation using ONNXRuntime + FastAPI (CPU-only)
• Build a Docker image for the runtime (mount the model folder at runtime; avoid baking multi-GB artifacts into the image)

Important notes

• Exporting ONNX models requires substantial RAM (12–32 GB recommended for the v1.5 pipeline).
• Export/quantize should be done on a machine with enough memory; serving can be done on smaller servers.
• Quantization reduces model size and improves CPU speed but may slightly affect output quality.
• This example is for demo and experimentation. For production you must add rate limiting, job queueing, concurrency controls, and safety filtering.

Quickstart (local export, quantize, run)

1. Prepare environment (recommended on an export-capable machine) python -m venv venv source venv/bin/activate pip install --upgrade pip

Optionally install a CPU PyTorch wheel first:

pip install --index-url https://download.pytorch.org/whl/cpu torch

pip install -r requirements.txt

2. Export ONNX models (this will download the HF model and export three ONNX files) python export_to_onnx.py --model_id runwayml/stable-diffusion-v1-5 --out_dir onnx_models

3. Quantize ONNX models python quantize_onnx.py --in_dir onnx_models --out_dir onnx_models_quant

4. Run the FastAPI server locally

Ensure onnx_models_quant folder is present

uvicorn server:app --host 0.0.0.0 --port 8000

5. Request an image POST to http://localhost:8000/generate with JSON body: {"prompt": "a fantasy landscape, colorful", "steps": 20, "scale": 7.5}

The server returns a JSON object with "image_base64" and "elapsed_s".

Docker (serve on a CPU host)

• Don't bake the model into the image. Mount it at runtime.

Build: docker build -t sd-onnx-cpu .

Run (mount models from host): docker run --rm -p 8000:8000
-v /absolute/path/onnx_models_quant:/app/onnx_models_quant
--cpus=4 --memory=16g
sd-onnx-cpu

Environment tuning

• Set INTRA_THREADS env var to control intra-op threads for ONNXRuntime: INTRA_THREADS=8 uvicorn server:app --host 0.0.0.0 --port 8000
• For best throughput, tune intra_op_num_threads close to the number of physical cores you dedicate to the container.

Troubleshooting

• If onnx export fails due to module forward signatures, the diffusers version may differ — please paste the error and I'll adapt the exporter to your diffusers/torch versions.
• If tokenizer/encoder shapes mismatch, ensure the same tokenizer model is used when exporting and when serving. If your chosen HF model uses a particular tokenizer, load that tokenizer in server.py.

Next steps / improvements

• Use Hugging Face Optimum or diffusers' built-in ONNX exporter for a simpler path if your versions align.
• Use onnxruntime with OpenVINO or oneDNN execution provider for additional CPU gains.
• Implement static calibration quantization (better INT8 results) using onnxruntime tools.
• Add a background job queue (RQ/Celery) and a small frontend for async UX.

If you want, I can:

• Generate this as a zip/archive for direct download.
• Add an Optimum-based exporter variant (shorter, often more robust).
• Add instructions to automatically download the exported models at container start (so images remain small).
• Help tune settings for a specific CPU (tell me CPU model and RAM).