Initial AuroraIMG-6M API

An authenticated CPU text-to-image API for direct-img, serving an ONNX
conversion of AuroraAI-Research/AuroraIMG-6M (Apache-2.0). About 0.4s per
image and ~150-200 MB RAM. Output matches the original pipeline to within
1/255 per pixel.

Pushes to main build ghcr.io/direct-img/auroraimg-6m:main and trigger the
Dokploy deploy webhook.

Co-Authored-By: Claude Opus 5.5 <noreply@anthropic.com>

.dockerignore

@@ -0,0 +1,6 @@
+.git
+.github
+.hf
+.venv
+scripts
+*.md

.github/workflows/docker.yml

@@ -0,0 +1,45 @@
+name: Build and Deploy
+
+on:
+  push:
+    branches: ['main']
+  workflow_dispatch:
+
+concurrency:
+  group: ${{ github.workflow }}-${{ github.ref }}
+  cancel-in-progress: true
+
+env:
+  IMAGE: ghcr.io/direct-img/auroraimg-6m
+
+jobs:
+  build-and-deploy:
+    runs-on: ubuntu-latest
+    permissions:
+      contents: read
+      packages: write
+
+    steps:
+      - uses: actions/checkout@v4
+
+      - uses: docker/setup-buildx-action@v3
+
+      - uses: docker/login-action@v3
+        with:
+          registry: ghcr.io
+          username: ${{ github.actor }}
+          password: ${{ secrets.GITHUB_TOKEN }}
+
+      - uses: docker/build-push-action@v6
+        with:
+          context: .
+          push: true
+          tags: |
+            ${{ env.IMAGE }}:main
+            ${{ env.IMAGE }}:sha-${{ github.sha }}
+          labels: org.opencontainers.image.source=${{ github.server_url }}/${{ github.repository }}
+          cache-from: type=gha
+          cache-to: type=gha,mode=max
+
+      - name: Trigger Dokploy deployment
+        run: curl -fsS -X POST "${{ secrets.DOKPLOY_WEBHOOK_URL }}"

.gitignore

@@ -0,0 +1,3 @@
+.hf/
+.venv/
+__pycache__/

Dockerfile

@@ -0,0 +1,13 @@
+FROM python:3.12-slim
+COPY --from=ghcr.io/astral-sh/uv:0.9 /uv /usr/local/bin/uv
+WORKDIR /app
+COPY requirements.txt .
+RUN uv pip install --system --no-cache -r requirements.txt
+COPY model model
+COPY app.py .
+RUN useradd --system --uid 10001 app
+USER app
+ENV PORT=8080 PYTHONUNBUFFERED=1
+EXPOSE 8080
+HEALTHCHECK --interval=30s --timeout=5s --start-period=15s CMD python -c "import os, urllib.request; urllib.request.urlopen(f'http://127.0.0.1:{os.environ[\"PORT\"]}/health', timeout=4)"
+CMD ["python", "app.py"]

README.md

@@ -0,0 +1,55 @@
+# AuroraIMG-6M API
+
+A small, authenticated text-to-image API for direct-img. It runs [AuroraIMG-6M](https://huggingface.co/AuroraAI-Research/AuroraIMG-6M), a 6.8M-parameter model, on CPU with onnxruntime. An image takes about 0.4s and the container uses about 150-200 MB of RAM.
+
+The output is 64×64 and upscaled, so images are soft and painterly: right colors, setting and mood, often not the exact object. That suits an abstract fallback when no real image exists.
+
+## API
+
+```
+GET /generate?prompt=frog+wizard+casting+a+spell
+Authorization: Bearer $AURORA_API_KEY
+```
+
+| Param | Default | Range |
+|---|---|---|
+| `prompt` | required | 1-200 characters |
+| `size` | `512` | 64-1024 (output is square, upscaled from 64×64) |
+| `format` | `webp` | `webp`, `png`, `jpeg` |
+| `steps` | `12` | 1-50 (12 looks the same as 50, at a quarter of the time) |
+| `cfg` | `6` | 0-20 (guidance; ≤1 disables it) |
+| `seed` | CRC32 of the lowercased prompt | 0-4294967295 |
+
+The same prompt and params always give the same image. Responses include `X-Seed` and `X-Generation-Ms`.
+
+Errors are JSON `{"error": "..."}`: 400 bad params, 401 missing or wrong key, 503 when the queue is full (`Retry-After: 2`).
+
+`GET /health` returns `ok` without auth.
+
+## Config
+
+| Env var | Default | |
+|---|---|---|
+| `AURORA_API_KEY` | none | Required. Without it `/generate` rejects every request |
+| `PORT` | `8080` | |
+| `THREADS` | `1` | onnxruntime threads per generation |
+| `CONCURRENCY` | `2` | Generations at once; more requests queue |
+| `MAX_QUEUE` | `32` | Queued requests beyond this get 503 |
+
+## Deploy
+
+Pushing to `main` builds the image with GitHub Actions, pushes it to `ghcr.io/direct-img/auroraimg-6m:main`, then calls the Dokploy deploy webhook (repo secret `DOKPLOY_WEBHOOK_URL`). It runs as the AuroraIMG-6M app in the direct-img Dokploy project.
+
+Run locally:
+
+```sh
+pip install -r requirements.txt
+AURORA_API_KEY=dev python app.py
+curl -H "Authorization: Bearer dev" "localhost:8080/generate?prompt=orange+cat" -o cat.webp
+```
+
+## Model and license
+
+`model/` holds an ONNX conversion of [AuroraAI-Research/AuroraIMG-6M](https://huggingface.co/AuroraAI-Research/AuroraIMG-6M) at revision `3d66743`, licensed under the Apache License 2.0 (copy in [`model/LICENSE`](model/LICENSE)). The weights are unchanged; they were converted from the original safetensors to ONNX, and the latent un-normalization and output scaling were folded into `decoder.onnx`. Output matches the original pipeline to within 1/255 per pixel.
+
+`scripts/export_onnx.py` regenerates `model/` from the original repo (needs torch, onnx, onnxscript, onnxruntime, safetensors).

app.py

@@ -0,0 +1,104 @@
+# AuroraIMG-6M image API. GET /generate?prompt=... with "Authorization: Bearer $AURORA_API_KEY" returns an image.
+# Runs the ONNX export of https://huggingface.co/AuroraAI-Research/AuroraIMG-6M on CPU with onnxruntime.
+import os, re, io, json, hmac, time, zlib, threading
+from http.server import ThreadingHTTPServer, BaseHTTPRequestHandler
+from urllib.parse import urlparse, parse_qs
+from pathlib import Path
+import numpy as np, onnxruntime as ort
+from PIL import Image
+from tokenizers import Tokenizer
+
+KEY = os.environ.get("AURORA_API_KEY", "").encode()
+PORT = int(os.environ.get("PORT", 8080))
+THREADS = int(os.environ.get("THREADS", 1))          # onnxruntime threads per generation
+CONCURRENCY = int(os.environ.get("CONCURRENCY", 2))  # generations at once; the rest queue
+MAX_QUEUE = int(os.environ.get("MAX_QUEUE", 32))     # queued requests beyond this get 503
+FORMATS = {"webp": ("WEBP", "image/webp", {"quality": 90}), "png": ("PNG", "image/png", {}), "jpeg": ("JPEG", "image/jpeg", {"quality": 90})}
+
+M = Path(__file__).parent / "model"
+opts = ort.SessionOptions()
+opts.intra_op_num_threads, opts.inter_op_num_threads = THREADS, 1
+S = {n: ort.InferenceSession(M / f"{n}.onnx", opts, providers=["CPUExecutionProvider"]) for n in ("text_encoder", "denoiser", "decoder")}
+run = lambda n, **x: S[n].run(None, x)[0]
+tok = Tokenizer.from_file(str(M / "tokenizer.json"))
+
+# Same cleanup and layout as the reference pipeline: <bos>=1, up to 62 tokens, <eos>=2, padded with 0 to 64
+def encode(p):
+    p = re.sub(r"https?://\S+|www\.\S+", " ", re.sub(r"<[^>]+>", " ", p))
+    seq = [1] + tok.encode(re.sub(r"\s+", " ", p.lower()).strip(), add_special_tokens=False).ids[:62] + [2]
+    t = np.array([seq + [0] * (64 - len(seq))], np.int64)
+    return t, t != 0
+
+NULL_TOK, NULL_MASK = encode("")
+NULL_CTX = run("text_encoder", tok=NULL_TOK, mask=NULL_MASK)
+
+def generate(prompt, steps, cfg, seed):
+    t, m = encode(prompt)
+    ctx = run("text_encoder", tok=t, mask=m)
+    z = np.random.default_rng(seed).standard_normal((1, 8, 16, 16), dtype=np.float32)
+    ts = np.linspace(1, 0, steps + 1, dtype=np.float32)
+    for a, b in zip(ts[:-1], ts[1:]):
+        if cfg > 1:
+            v_c, v_u = run("denoiser", x=np.concatenate([z, z]), t=np.full(2, a * 1000, np.float32),
+                           ctx=np.concatenate([ctx, NULL_CTX]), mask=np.concatenate([m, NULL_MASK]))
+            v = v_u + cfg * (v_c - v_u)
+            v *= np.linalg.norm(v_c) / max(np.linalg.norm(v), 1e-8)  # the pipeline's default cfg_rescale=1.0
+        else:
+            v = run("denoiser", x=z, t=np.full(1, a * 1000, np.float32), ctx=ctx, mask=m)
+        z = z + (b - a) * v
+    img = run("decoder", z=z)[0].transpose(1, 2, 0)
+    return Image.fromarray((img * 255).round().astype(np.uint8))
+
+def params(q):
+    def num(k, cast, lo, hi, default):
+        try: v = cast(q[k]) if k in q else default
+        except ValueError: raise ValueError(f"{k} must be a number")
+        if not lo <= v <= hi: raise ValueError(f"{k} must be {lo}-{hi}")
+        return v
+    prompt = q.get("prompt", "").strip()
+    if not 1 <= len(prompt) <= 200: raise ValueError("prompt must be 1-200 characters")
+    fmt = q.get("format", "webp").lower()
+    if fmt not in FORMATS: raise ValueError(f"format must be one of {', '.join(FORMATS)}")
+    return dict(prompt=prompt, steps=num("steps", int, 1, 50, 12), cfg=num("cfg", float, 0, 20, 6.0),
+                seed=num("seed", int, 0, 2**32 - 1, zlib.crc32(prompt.lower().encode())), size=num("size", int, 64, 1024, 512), fmt=fmt)
+
+gen_slots, queue_lock, queued = threading.BoundedSemaphore(CONCURRENCY), threading.Lock(), 0
+
+class Handler(BaseHTTPRequestHandler):
+    protocol_version = "HTTP/1.1"
+
+    def reply(s, code, body, ctype="application/json", **headers):
+        s.send_response(code)
+        for k, v in {"Content-Type": ctype, "Content-Length": len(body), "Cache-Control": "no-store", **headers}.items(): s.send_header(k, str(v))
+        s.end_headers(); s.wfile.write(body)
+
+    err = lambda s, code, msg, **h: s.reply(code, json.dumps({"error": msg}).encode(), **h)
+
+    def do_GET(s):
+        global queued
+        u = urlparse(s.path)
+        if u.path == "/health": return s.reply(200, b"ok", "text/plain")
+        if u.path != "/generate": return s.err(404, "not found")
+        if not KEY or not hmac.compare_digest(s.headers.get("Authorization", "").encode(), b"Bearer " + KEY): return s.err(401, "unauthorized")
+        try: p = params({k: v[-1] for k, v in parse_qs(u.query).items()})
+        except ValueError as e: return s.err(400, str(e))
+        with queue_lock:
+            if queued >= CONCURRENCY + MAX_QUEUE: return s.err(503, "busy, retry shortly", **{"Retry-After": 2})
+            queued += 1
+        try:
+            t0 = time.perf_counter()
+            with gen_slots: img = generate(p["prompt"], p["steps"], p["cfg"], p["seed"])
+            if p["size"] != 64: img = img.resize((p["size"],) * 2, Image.LANCZOS)
+            fmt, ctype, save = FORMATS[p["fmt"]]
+            buf = io.BytesIO(); img.save(buf, fmt, **save)
+            s.reply(200, buf.getvalue(), ctype, **{"X-Seed": p["seed"], "X-Generation-Ms": round((time.perf_counter() - t0) * 1000)})
+        finally:
+            with queue_lock: queued -= 1
+
+if __name__ == "__main__":
+    if not KEY: print("warning: AURORA_API_KEY is not set, so /generate rejects every request", flush=True)
+    generate("warmup", 2, 6.0, 0)
+    server = ThreadingHTTPServer(("0.0.0.0", PORT), Handler)
+    server.daemon_threads = True
+    print(f"listening on :{PORT} (threads={THREADS}, concurrency={CONCURRENCY})", flush=True)
+    server.serve_forever()

model/LICENSE

@@ -0,0 +1,202 @@
+
+                                 Apache License
+                           Version 2.0, January 2004
+                        http://www.apache.org/licenses/
+
+   TERMS AND CONDITIONS FOR USE, REPRODUCTION, AND DISTRIBUTION
+
+   1. Definitions.
+
+      "License" shall mean the terms and conditions for use, reproduction,
+      and distribution as defined by Sections 1 through 9 of this document.
+
+      "Licensor" shall mean the copyright owner or entity authorized by
+      the copyright owner that is granting the License.
+
+      "Legal Entity" shall mean the union of the acting entity and all
+      other entities that control, are controlled by, or are under common
+      control with that entity. For the purposes of this definition,
+      "control" means (i) the power, direct or indirect, to cause the
+      direction or management of such entity, whether by contract or
+      otherwise, or (ii) ownership of fifty percent (50%) or more of the
+      outstanding shares, or (iii) beneficial ownership of such entity.
+
+      "You" (or "Your") shall mean an individual or Legal Entity
+      exercising permissions granted by this License.
+
+      "Source" form shall mean the preferred form for making modifications,
+      including but not limited to software source code, documentation
+      source, and configuration files.
+
+      "Object" form shall mean any form resulting from mechanical
+      transformation or translation of a Source form, including but
+      not limited to compiled object code, generated documentation,
+      and conversions to other media types.
+
+      "Work" shall mean the work of authorship, whether in Source or
+      Object form, made available under the License, as indicated by a
+      copyright notice that is included in or attached to the work
+      (an example is provided in the Appendix below).
+
+      "Derivative Works" shall mean any work, whether in Source or Object
+      form, that is based on (or derived from) the Work and for which the
+      editorial revisions, annotations, elaborations, or other modifications
+      represent, as a whole, an original work of authorship. For the purposes
+      of this License, Derivative Works shall not include works that remain
+      separable from, or merely link (or bind by name) to the interfaces of,
+      the Work and Derivative Works thereof.
+
+      "Contribution" shall mean any work of authorship, including
+      the original version of the Work and any modifications or additions
+      to that Work or Derivative Works thereof, that is intentionally
+      submitted to Licensor for inclusion in the Work by the copyright owner
+      or by an individual or Legal Entity authorized to submit on behalf of
+      the copyright owner. For the purposes of this definition, "submitted"
+      means any form of electronic, verbal, or written communication sent
+      to the Licensor or its representatives, including but not limited to
+      communication on electronic mailing lists, source code control systems,
+      and issue tracking systems that are managed by, or on behalf of, the
+      Licensor for the purpose of discussing and improving the Work, but
+      excluding communication that is conspicuously marked or otherwise
+      designated in writing by the copyright owner as "Not a Contribution."
+
+      "Contributor" shall mean Licensor and any individual or Legal Entity
+      on behalf of whom a Contribution has been received by Licensor and
+      subsequently incorporated within the Work.
+
+   2. Grant of Copyright License. Subject to the terms and conditions of
+      this License, each Contributor hereby grants to You a perpetual,
+      worldwide, non-exclusive, no-charge, royalty-free, irrevocable
+      copyright license to reproduce, prepare Derivative Works of,
+      publicly display, publicly perform, sublicense, and distribute the
+      Work and such Derivative Works in Source or Object form.
+
+   3. Grant of Patent License. Subject to the terms and conditions of
+      this License, each Contributor hereby grants to You a perpetual,
+      worldwide, non-exclusive, no-charge, royalty-free, irrevocable
+      (except as stated in this section) patent license to make, have made,
+      use, offer to sell, sell, import, and otherwise transfer the Work,
+      where such license applies only to those patent claims licensable
+      by such Contributor that are necessarily infringed by their
+      Contribution(s) alone or by combination of their Contribution(s)
+      with the Work to which such Contribution(s) was submitted. If You
+      institute patent litigation against any entity (including a
+      cross-claim or counterclaim in a lawsuit) alleging that the Work
+      or a Contribution incorporated within the Work constitutes direct
+      or contributory patent infringement, then any patent licenses
+      granted to You under this License for that Work shall terminate
+      as of the date such litigation is filed.
+
+   4. Redistribution. You may reproduce and distribute copies of the
+      Work or Derivative Works thereof in any medium, with or without
+      modifications, and in Source or Object form, provided that You
+      meet the following conditions:
+
+      (a) You must give any other recipients of the Work or
+          Derivative Works a copy of this License; and
+
+      (b) You must cause any modified files to carry prominent notices
+          stating that You changed the files; and
+
+      (c) You must retain, in the Source form of any Derivative Works
+          that You distribute, all copyright, patent, trademark, and
+          attribution notices from the Source form of the Work,
+          excluding those notices that do not pertain to any part of
+          the Derivative Works; and
+
+      (d) If the Work includes a "NOTICE" text file as part of its
+          distribution, then any Derivative Works that You distribute must
+          include a readable copy of the attribution notices contained
+          within such NOTICE file, excluding those notices that do not
+          pertain to any part of the Derivative Works, in at least one
+          of the following places: within a NOTICE text file distributed
+          as part of the Derivative Works; within the Source form or
+          documentation, if provided along with the Derivative Works; or,
+          within a display generated by the Derivative Works, if and
+          wherever such third-party notices normally appear. The contents
+          of the NOTICE file are for informational purposes only and
+          do not modify the License. You may add Your own attribution
+          notices within Derivative Works that You distribute, alongside
+          or as an addendum to the NOTICE text from the Work, provided
+          that such additional attribution notices cannot be construed
+          as modifying the License.
+
+      You may add Your own copyright statement to Your modifications and
+      may provide additional or different license terms and conditions
+      for use, reproduction, or distribution of Your modifications, or
+      for any such Derivative Works as a whole, provided Your use,
+      reproduction, and distribution of the Work otherwise complies with
+      the conditions stated in this License.
+
+   5. Submission of Contributions. Unless You explicitly state otherwise,
+      any Contribution intentionally submitted for inclusion in the Work
+      by You to the Licensor shall be under the terms and conditions of
+      this License, without any additional terms or conditions.
+      Notwithstanding the above, nothing herein shall supersede or modify
+      the terms of any separate license agreement you may have executed
+      with Licensor regarding such Contributions.
+
+   6. Trademarks. This License does not grant permission to use the trade
+      names, trademarks, service marks, or product names of the Licensor,
+      except as required for reasonable and customary use in describing the
+      origin of the Work and reproducing the content of the NOTICE file.
+
+   7. Disclaimer of Warranty. Unless required by applicable law or
+      agreed to in writing, Licensor provides the Work (and each
+      Contributor provides its Contributions) on an "AS IS" BASIS,
+      WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or
+      implied, including, without limitation, any warranties or conditions
+      of TITLE, NON-INFRINGEMENT, MERCHANTABILITY, or FITNESS FOR A
+      PARTICULAR PURPOSE. You are solely responsible for determining the
+      appropriateness of using or redistributing the Work and assume any
+      risks associated with Your exercise of permissions under this License.
+
+   8. Limitation of Liability. In no event and under no legal theory,
+      whether in tort (including negligence), contract, or otherwise,
+      unless required by applicable law (such as deliberate and grossly
+      negligent acts) or agreed to in writing, shall any Contributor be
+      liable to You for damages, including any direct, indirect, special,
+      incidental, or consequential damages of any character arising as a
+      result of this License or out of the use or inability to use the
+      Work (including but not limited to damages for loss of goodwill,
+      work stoppage, computer failure or malfunction, or any and all
+      other commercial damages or losses), even if such Contributor
+      has been advised of the possibility of such damages.
+
+   9. Accepting Warranty or Additional Liability. While redistributing
+      the Work or Derivative Works thereof, You may choose to offer,
+      and charge a fee for, acceptance of support, warranty, indemnity,
+      or other liability obligations and/or rights consistent with this
+      License. However, in accepting such obligations, You may act only
+      on Your own behalf and on Your sole responsibility, not on behalf
+      of any other Contributor, and only if You agree to indemnify,
+      defend, and hold each Contributor harmless for any liability
+      incurred by, or claims asserted against, such Contributor by reason
+      of your accepting any such warranty or additional liability.
+
+   END OF TERMS AND CONDITIONS
+
+   APPENDIX: How to apply the Apache License to your work.
+
+      To apply the Apache License to your work, attach the following
+      boilerplate notice, with the fields enclosed by brackets "[]"
+      replaced with your own identifying information. (Don't include
+      the brackets!)  The text should be enclosed in the appropriate
+      comment syntax for the file format. We also recommend that a
+      file or class name and description of purpose be included on the
+      same "printed page" as the copyright notice for easier
+      identification within third-party archives.
+
+   Copyright [yyyy] [name of copyright owner]
+
+   Licensed under the Apache License, Version 2.0 (the "License");
+   you may not use this file except in compliance with the License.
+   You may obtain a copy of the License at
+
+       http://www.apache.org/licenses/LICENSE-2.0
+
+   Unless required by applicable law or agreed to in writing, software
+   distributed under the License is distributed on an "AS IS" BASIS,
+   WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
+   See the License for the specific language governing permissions and
+   limitations under the License.

requirements.txt

@@ -0,0 +1,4 @@
+onnxruntime==1.31.0
+numpy==2.5.4
+pillow==12.3.0
+tokenizers==0.23.3

scripts/export_onnx.py

@@ -0,0 +1,63 @@
+# Converts AuroraIMG-6M (https://huggingface.co/AuroraAI-Research/AuroraIMG-6M, Apache-2.0) to the ONNX files in model/.
+# Only needed to regenerate them. Needs: torch onnx onnxscript onnxruntime safetensors tokenizers numpy
+#   python scripts/export_onnx.py
+# Weights are read from the .safetensors files only; the repo's pickle .pt files are never downloaded.
+import sys, types, json, shutil, urllib.request, importlib.util
+from pathlib import Path
+import numpy as np, torch, onnxruntime as ort
+from safetensors.torch import load_file
+torch.set_grad_enabled(False)
+
+REPO, REV = "AuroraAI-Research/AuroraIMG-6M", "3d66743c3139d034b832a670d067221f5d544bbb"
+FILES = ["pipeline.py", "tokenizer/tokenizer.json", "vae/config.json"] + [f"{m}/diffusion_pytorch_model.safetensors" for m in ("text_encoder", "unet", "vae")]
+ROOT = Path(__file__).resolve().parent.parent
+HF, OUT = ROOT / ".hf", ROOT / "model"
+
+for f in FILES:
+    p = HF / f
+    if not p.exists():
+        p.parent.mkdir(parents=True, exist_ok=True)
+        urllib.request.urlretrieve(f"https://huggingface.co/{REPO}/resolve/{REV}/{f}", p)
+
+# pipeline.py imports diffusers/transformers only for base classes; these stubs let it load with plain torch
+class ConfigMixin:
+    def register_to_config(self, **kw): self.__dict__["config"] = types.SimpleNamespace(**kw)
+stub = lambda name, **attrs: types.SimpleNamespace(__name__=name, **attrs)
+sys.modules |= {"diffusers": stub("diffusers", DiffusionPipeline=object, ModelMixin=type("ModelMixin", (), {}), ConfigMixin=ConfigMixin),
+                "diffusers.utils": stub("diffusers.utils", BaseOutput=object), "transformers": stub("transformers", PreTrainedTokenizerFast=object)}
+spec = importlib.util.spec_from_file_location("aurora_pipeline", HF / "pipeline.py")
+P = importlib.util.module_from_spec(spec); spec.loader.exec_module(P)
+
+def load(m, sub):
+    m.load_state_dict(load_file(HF / sub / "diffusion_pytorch_model.safetensors"), strict=True); return m.eval()
+vae_cfg = json.loads((HF / "vae/config.json").read_text())
+te, den = load(P.TextEncoder(vocab=4096, d=128, layers=4), "text_encoder"), load(P.Denoiser(), "unet")
+vae = load(P.VAE(lat_mu=vae_cfg["lat_mu"], lat_sd=vae_cfg["lat_sd"]), "vae")
+
+class Decode(torch.nn.Module):
+    """Un-normalizes latents and maps the VAE output to 0..1 RGB, as the original pipeline does after sampling."""
+    def __init__(s):
+        super().__init__(); s.vae = vae
+        s.register_buffer("sd", torch.tensor(vae_cfg["lat_sd"]).view(1, 8, 1, 1)); s.register_buffer("mu", torch.tensor(vae_cfg["lat_mu"]).view(1, 8, 1, 1))
+    def forward(s, z): return (s.vae.decode(z * s.sd + s.mu) / 2 + 0.5).clamp(0, 1)
+dec = Decode().eval()
+
+B = torch.export.Dim("batch", min=1, max=64)
+tok = torch.randint(3, 4096, (2, 64)); tok[:, 10:] = 0; tok[:, 0] = 1; mask = tok != 0
+ctx, x, t = te(tok, mask), torch.randn(2, 8, 16, 16), torch.tensor([500.0, 500.0])
+OUT.mkdir(exist_ok=True)
+graphs = {
+    "text_encoder": (te, (tok, mask), ["tok", "mask"], ["ctx"]),
+    "denoiser": (den, (x, t, ctx, mask), ["x", "t", "ctx", "mask"], ["v"]),
+    "decoder": (dec, (x,), ["z"], ["image"]),
+}
+for name, (m, args, ins, outs) in graphs.items():
+    shapes = tuple({0: B} for _ in args)
+    torch.onnx.export(m, args, OUT / f"{name}.onnx", input_names=ins, output_names=outs, dynamic_shapes=shapes, dynamo=True, external_data=False)
+    sess = ort.InferenceSession(OUT / f"{name}.onnx", providers=["CPUExecutionProvider"])
+    want = m(*args).numpy()
+    got = sess.run(None, {k: a.numpy() for k, a in zip(ins, args)})[0]
+    err = float(np.abs(got - want).max())
+    print(f"{name}.onnx: {(OUT / f'{name}.onnx').stat().st_size / 1e6:.1f} MB, max abs diff vs torch {err:.2e}")
+    assert err < 1e-3, name
+shutil.copyfile(HF / "tokenizer/tokenizer.json", OUT / "tokenizer.json")