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>

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")