mirror of
https://github.com/direct-img/AuroraIMG-6M.git
synced 2026-10-29 09:04:52 +00:00
multipleof4andClaude Opus 5.5
efe9971809
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>
64 lines
3.7 KiB
Python
64 lines
3.7 KiB
Python
# 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")
|