Initial release

.gitattributes

@@ -0,0 +1,35 @@
+*.7z filter=lfs diff=lfs merge=lfs -text
+*.arrow filter=lfs diff=lfs merge=lfs -text
+*.bin filter=lfs diff=lfs merge=lfs -text
+*.bz2 filter=lfs diff=lfs merge=lfs -text
+*.ckpt filter=lfs diff=lfs merge=lfs -text
+*.ftz filter=lfs diff=lfs merge=lfs -text
+*.gz filter=lfs diff=lfs merge=lfs -text
+*.h5 filter=lfs diff=lfs merge=lfs -text
+*.joblib filter=lfs diff=lfs merge=lfs -text
+*.lfs.* filter=lfs diff=lfs merge=lfs -text
+*.mlmodel filter=lfs diff=lfs merge=lfs -text
+*.model filter=lfs diff=lfs merge=lfs -text
+*.msgpack filter=lfs diff=lfs merge=lfs -text
+*.npy filter=lfs diff=lfs merge=lfs -text
+*.npz filter=lfs diff=lfs merge=lfs -text
+*.onnx filter=lfs diff=lfs merge=lfs -text
+*.ot filter=lfs diff=lfs merge=lfs -text
+*.parquet filter=lfs diff=lfs merge=lfs -text
+*.pb filter=lfs diff=lfs merge=lfs -text
+*.pickle filter=lfs diff=lfs merge=lfs -text
+*.pkl filter=lfs diff=lfs merge=lfs -text
+*.pt filter=lfs diff=lfs merge=lfs -text
+*.pth filter=lfs diff=lfs merge=lfs -text
+*.rar filter=lfs diff=lfs merge=lfs -text
+*.safetensors filter=lfs diff=lfs merge=lfs -text
+saved_model/**/* filter=lfs diff=lfs merge=lfs -text
+*.tar.* filter=lfs diff=lfs merge=lfs -text
+*.tar filter=lfs diff=lfs merge=lfs -text
+*.tflite filter=lfs diff=lfs merge=lfs -text
+*.tgz filter=lfs diff=lfs merge=lfs -text
+*.wasm filter=lfs diff=lfs merge=lfs -text
+*.xz filter=lfs diff=lfs merge=lfs -text
+*.zip filter=lfs diff=lfs merge=lfs -text
+*.zst filter=lfs diff=lfs merge=lfs -text
+*tfevents* filter=lfs diff=lfs merge=lfs -text

Nucleus-Image-FP8.safetensors

@@ -0,0 +1,3 @@
+version https://git-lfs.github.com/spec/v1
+oid sha256:6555d4212cf6eb00d0f03505cc8f3e2032ca73927905cd57b54df2eb999115f7
+size 16942185128

README.md

@@ -0,0 +1,224 @@
+---
+license: apache-2.0
+language:
+- en
+library_name: diffusers
+pipeline_tag: text-to-image
+base_model: NucleusAI/Nucleus-Image
+tags:
+- moe
+- sparse-moe
+- diffusion
+- text-to-image
+- image-generation
+- quantization
+- fp8
+---
+
+# Nucleus-Image — FP8 (e4m3)
+
+FP8 weight-only quantization of [`NucleusAI/Nucleus-Image`](https://huggingface.co/NucleusAI/Nucleus-Image). Single 16.94 GB safetensors file at the repo root (was 33.85 GB in BF16). Peak VRAM at 1024² is **17.6 GB** with `enable_model_cpu_offload()` — fits comfortably on 24 GB cards.
+
+## Run it
+
+```python
+import importlib.util, torch
+from diffusers import DiffusionPipeline
+from huggingface_hub import hf_hub_download
+
+REPO = "D-Squarius-Green-Jr/Nucleus-Image-FP8"
+patch_py = hf_hub_download(REPO, "moe_fp8_patch.py")
+weights  = hf_hub_download(REPO, "Nucleus-Image-FP8.safetensors")
+hf_hub_download(REPO, "config.json")  # sits next to the weights
+
+spec = importlib.util.spec_from_file_location("moe_fp8_patch", patch_py)
+patch = importlib.util.module_from_spec(spec); spec.loader.exec_module(patch)
+patch.apply_patch()
+
+transformer = patch.load_fp8_safetensors_transformer(weights)
+
+pipe = DiffusionPipeline.from_pretrained(
+    "NucleusAI/Nucleus-Image",
+    transformer=transformer,
+    torch_dtype=torch.bfloat16,
+)
+pipe.enable_model_cpu_offload()
+
+image = pipe(
+    prompt="A quiet alpine lake at sunrise, mist rising off still water, snow-capped peaks reflected, soft pink and gold sky",
+    width=1024, height=1024,
+    num_inference_steps=20, guidance_scale=8.0,
+    generator=torch.Generator(device="cuda").manual_seed(42),
+).images[0]
+image.save("out.png")
+```
+
+> Skip `TextKVCacheConfig` if running multiple prompts of different lengths — its state has no public reset (as of `diffusers` 0.38.0.dev0). Single prompts are fine.
+
+## Quant scheme
+
+| Layers | How | Size |
+|---|---|---|
+| 354 `nn.Linear` | FP8 e4m3, per-output-channel scale | 1.6 GB |
+| 29 `SwiGLUExperts` (15.3 B params, fused 64-expert tensors) | FP8 e4m3, per-expert × per-output-channel scale | 15.3 GB |
+| Norms / routing logits / embeddings | BF16 (routing decisions stay bit-identical to BF16) | 0.04 GB |
+
+Off-the-shelf quantizers (TorchAO, optimum-quanto, bnb) only walk `nn.Linear`, missing the 15.3 B inside `SwiGLUExperts`. The runtime patch handles both.
+
+## Numbers
+
+- Dequant rel L2: 2.69 % worst across all SwiGLUExperts
+- Per-layer forward rel L2 (random `x`): 4.49 %
+- Generation: ~25-42 s per image at 1024², 20 steps, RTX 5090
+- Peak VRAM: 17.62 GB allocated / 17.82 GB reserved
+
+---
+
+# Original Model Card
+
+The text below is the verbatim model card from [`NucleusAI/Nucleus-Image`](https://huggingface.co/NucleusAI/Nucleus-Image).
+
+<p align="center"> <a href="https://withnucleus.ai/image" target="_blank" rel="noopener noreferrer"><img src="https://storage.googleapis.com/nucleus_image_v1/nucleus_header.png" width="400"/></a></p>
+<p align="center">
+    🌐 <a href="https://withnucleus.ai/image"><b>Website</b></a>&nbsp;&nbsp; | &nbsp;&nbsp;🖥️ <a href="https://github.com/WithNucleusAI/Nucleus-Image"><b>GitHub</b></a>&nbsp;&nbsp; | &nbsp;&nbsp;🤗 <a href="https://huggingface.co/NucleusAI/NucleusMoE-Image"><b>Hugging Face</b></a>&nbsp;&nbsp; | &nbsp;&nbsp;📑 <a href="https://storage.googleapis.com/nucleus_image_v1/Nucleus-Image-Technical-Report.pdf"><b>Tech Report</b></a>
+</p>
+
+## Introduction
+
+**Nucleus-Image** is a text-to-image generation model built on a sparse mixture-of-experts (MoE) diffusion transformer architecture. It scales to **17B total parameters** across 64 routed experts per layer while activating only **~2B parameters** per forward pass, establishing a new Pareto frontier in quality-versus-efficiency. Nucleus-Image matches or exceeds leading models including Qwen-Image, GPT Image 1, Seedream 3.0, and Imagen4 on GenEval, DPG-Bench, and OneIG-Bench. This is a **base model** released without any post-training optimization (no DPO, no reinforcement learning, no human preference tuning). All reported results reflect pre-training performance only. We release the full model weights, training code, and dataset, making Nucleus-Image the first fully open-source MoE diffusion model at this quality tier.
+
+## Key Features
+
+- **Sparse MoE efficiency**: 17B total capacity with only ~2B active parameters per forward pass, enabling high-quality generation at a fraction of the inference cost of dense models
+- **Expert-Choice Routing**: Guarantees balanced expert utilization without auxiliary load-balancing losses, with a decoupled routing design that separates timestep-aware assignment from timestep-conditioned computation
+- **Base model, no post-training**: This is a base model. All benchmark results are from pre-training alone, without DPO, reinforcement learning, or human preference tuning
+- **Multi-aspect-ratio support**: Trained with aspect-ratio bucketing from the outset at every resolution stage, supporting a range of output dimensions
+- **Text KV caching via diffusers**: Text tokens are excluded from the transformer backbone entirely and their KV projections are cached across all denoising steps. This caching is natively integrated into the `diffusers` pipeline. Simply enable it with `TextKVCacheConfig` for automatic speedup with no code changes to the inference loop
+- **Progressive resolution training**: Three-stage curriculum (256 → 512 → 1024) with progressive sparsification of expert capacity
+
+## Architecture
+
+![Architecture](https://storage.googleapis.com/nucleus_image_v1/Architecture_Diagram.png)
+
+Nucleus-Image is a 32-layer diffusion transformer where 29 of the 32 blocks replace the dense FFN with a sparse MoE layer containing 64 routed experts and one shared expert (the first 3 layers use dense FFN for training stability). Image queries attend to concatenated image and text key-value pairs via joint attention. Text tokens are excluded from the transformer backbone entirely, participating only as KV contributors. This eliminates MoE routing overhead for text and enables full text KV caching across denoising steps.
+
+Routing uses **Expert-Choice** with a **decoupled design**: the router receives the unmodulated token representation concatenated with the timestep embedding, while expert MLPs receive the fully modulated representation. This prevents the adaptive modulation scale — which varies by an order of magnitude across timesteps — from collapsing expert selection into timestep-dependent routing, preserving spatial and semantic expert specialization.
+
+## Model Specifications
+
+| Specification | Value |
+|---|---|
+| Total parameters | 17B |
+| Active parameters | ~2B |
+| Architecture | Sparse MoE Diffusion Transformer |
+| Layers | 32 |
+| Hidden dimension | 2048 |
+| Attention heads (Q / KV) | 16 / 4 (GQA) |
+| Experts per MoE layer | 64 routed + 1 shared |
+| Expert hidden dimension | 1344 |
+| Text encoder | Qwen3-VL-8B-Instruct |
+| Image tokenizer | Qwen-Image VAE (16ch) |
+| Training data | 700M images, 1.5B caption pairs |
+| Training curriculum | Progressive resolution (256 → 512 → 1024) |
+| Total training steps | 1.7M |
+
+## Benchmark Results
+
+![Overall Performance](https://storage.googleapis.com/nucleus_image_v1/Overall-Performance.png)
+
+Nucleus-Image achieves state-of-the-art or near state-of-the-art results on all three benchmarks despite activating only ~2B of its 17B parameters per forward pass. All results are from the base model at 1024x1024, 50 inference steps, CFG scale 8.0.
+
+| Benchmark | Score | Highlights |
+|---|---|---|
+| **GenEval** | **0.87** | Matches Qwen-Image; leads all models on spatial position (0.85) |
+| **DPG-Bench** | **88.79** | #1 overall; leads in entity (93.08), attribute (92.20), and other (93.62) |
+| **OneIG-Bench** | **0.522** | Surpasses Imagen4 (0.515) and Recraft V3 (0.502); strong style (0.430) |
+
+## Quick Start
+
+Install the latest version of diffusers:
+```
+pip install git+https://github.com/huggingface/diffusers
+```
+
+Generate images with Nucleus-Image:
+
+```python
+import torch
+from diffusers import DiffusionPipeline
+from diffusers import TextKVCacheConfig
+
+model_name = "NucleusAI/Nucleus-Image"
+
+pipe = DiffusionPipeline.from_pretrained(model_name, torch_dtype=torch.bfloat16)
+pipe.to("cuda")
+
+# Enable Text KV caching across denoising steps (integrated into diffusers)
+config = TextKVCacheConfig()
+pipe.transformer.enable_cache(config)
+
+# Supported aspect ratios
+aspect_ratios = {
+    "1:1": (1024, 1024),
+    "16:9": (1344, 768),
+    "9:16": (768, 1344),
+    "4:3": (1184, 896),
+    "3:4": (896, 1184),
+    "3:2": (1248, 832),
+    "2:3": (832, 1248),
+}
+
+prompt = "A weathered lighthouse on a rocky coastline at golden hour, waves crashing against the rocks below, seagulls circling overhead, dramatic clouds painted in shades of amber and violet"
+width, height = aspect_ratios["16:9"]
+
+image = pipe(
+    prompt=prompt,
+    width=width,
+    height=height,
+    num_inference_steps=50,
+    guidance_scale=8.0,
+    generator=torch.Generator(device="cuda").manual_seed(42),
+).images[0]
+
+image.save("nucleus_output.png")
+```
+
+## Highlights
+
+### Portraits & People
+
+Nucleus-Image generations of human subjects and portraits, spanning diverse cultures, ages, and artistic styles. From expressive character studies to fine-grained close-ups with intricate skin texture and detail.
+
+![](https://storage.googleapis.com/nucleus_image_v1/Collage-1-Top.jpeg)
+![](https://storage.googleapis.com/nucleus_image_v1/Collage-1-Bottom.jpeg)
+
+### Fantasy, Surrealism & Nature
+
+Nucleus-Image generations spanning fantasy, surrealism, animation, and the natural world.
+
+![](https://storage.googleapis.com/nucleus_image_v1/Collage-2-Top.jpeg)
+![](https://storage.googleapis.com/nucleus_image_v1/Collage-2-Bottom.jpeg)
+
+### Commercial & Everyday Imagery
+
+Nucleus-Image generations across product photography, architecture, typography, food, and world culture, demonstrating versatility in commercial, conceptual, and everyday imagery.
+
+![](https://storage.googleapis.com/nucleus_image_v1/Collage-3-Top.jpeg)
+![](https://storage.googleapis.com/nucleus_image_v1/Collage-3-Bottom.jpeg)
+
+## License
+
+Nucleus-Image is licensed under [Apache 2.0](https://www.apache.org/licenses/LICENSE-2.0).
+
+## Citation
+
+```bibtex
+@misc{nucleusimage2026,
+      title={Nucleus-Image: Sparse MoE for Image Generation},
+      author={Nucleus AI Team},
+      year={2026},
+      eprint={XXXX.XXXXX},
+      archivePrefix={arXiv},
+      primaryClass={cs.CV},
+}
+```

config.json

@@ -0,0 +1,60 @@
+{
+  "_class_name": "NucleusMoEImageTransformer2DModel",
+  "_diffusers_version": "0.38.0.dev0",
+  "_name_or_path": "C:\\Users\\gabeg\\Projects\\nucleus-image\\Nucleus-Image_FP8",
+  "attention_head_dim": 128,
+  "axes_dims_rope": [
+    16,
+    56,
+    56
+  ],
+  "capacity_factors": [
+    0.0,
+    0.0,
+    0.0,
+    4.0,
+    4.0,
+    2.0,
+    2.0,
+    2.0,
+    2.0,
+    2.0,
+    2.0,
+    2.0,
+    2.0,
+    2.0,
+    2.0,
+    2.0,
+    2.0,
+    2.0,
+    2.0,
+    2.0,
+    2.0,
+    2.0,
+    2.0,
+    2.0,
+    2.0,
+    2.0,
+    2.0,
+    2.0,
+    2.0,
+    2.0,
+    2.0,
+    2.0
+  ],
+  "dense_moe_strategy": "leave_first_three_blocks_dense",
+  "in_channels": 64,
+  "joint_attention_dim": 4096,
+  "mlp_ratio": 4.0,
+  "moe_enabled": true,
+  "moe_intermediate_dim": 1344,
+  "num_attention_heads": 16,
+  "num_experts": 64,
+  "num_key_value_heads": 4,
+  "num_layers": 32,
+  "out_channels": 16,
+  "patch_size": 2,
+  "route_scale": 2.5,
+  "use_grouped_mm": true,
+  "use_sigmoid": false
+}

moe_fp8_patch.py

@@ -0,0 +1,281 @@
+"""
+SwiGLUExperts FP8 monkey-patch.
+
+Imported by both 03_quantize_fp8.py (before loading BF16 model) and 04_test_inference.py
+(before loading FP8 model). Idempotent: importing twice is a no-op.
+
+What it does:
+- Adds two persistent buffers (gate_up_proj_scale, down_proj_scale) to every SwiGLUExperts.
+- Replaces _run_experts_for_loop to dequantize per-expert weights on-the-fly when stored as fp8_e4m3fn.
+- Forces use_grouped_mm=False (the grouped_mm kernel doesn't accept fp8 e4m3 inputs as of torch 2.11).
+- If weights are still bf16 (un-quantized model), behavior is identical to the original SwiGLU forward.
+"""
+from __future__ import annotations
+import torch
+import torch.nn as nn
+import torch.nn.functional as F
+from diffusers.models.transformers import transformer_nucleusmoe_image as _moe_mod
+
+FP8_E4M3_MAX = 448.0  # float8_e4m3fn dynamic range
+SCALE_DTYPE = torch.bfloat16
+
+_PATCH_FLAG = "_nucleus_fp8_patched_v1"
+
+
+def apply_patch():
+    if getattr(_moe_mod.SwiGLUExperts, _PATCH_FLAG, False):
+        return  # already patched
+
+    cls = _moe_mod.SwiGLUExperts
+    orig_init = cls.__init__
+
+    def patched_init(self, hidden_size, moe_intermediate_dim, num_experts, use_grouped_mm: bool = False):
+        # Force use_grouped_mm off — fp8 inputs aren't supported by F.grouped_mm on Blackwell yet,
+        # and the for-loop path is the patched/quantized path.
+        orig_init(self, hidden_size, moe_intermediate_dim, num_experts, use_grouped_mm=False)
+        # Persistent buffers; default to ones so a non-quantized BF16 checkpoint still produces
+        # mathematically identical output through patched_for_loop.
+        self.register_buffer(
+            "gate_up_proj_scale",
+            torch.ones(num_experts, 1, 2 * moe_intermediate_dim, dtype=SCALE_DTYPE),
+            persistent=True,
+        )
+        self.register_buffer(
+            "down_proj_scale",
+            torch.ones(num_experts, 1, hidden_size, dtype=SCALE_DTYPE),
+            persistent=True,
+        )
+
+    def patched_for_loop(self, x: torch.Tensor, num_tokens_per_expert: torch.Tensor) -> torch.Tensor:
+        n_list = num_tokens_per_expert.tolist()
+        n_real = sum(n_list)
+        n_pad = x.shape[0] - n_real
+        x_per_expert = torch.split(x[:n_real], split_size_or_sections=n_list, dim=0)
+
+        is_fp8 = self.gate_up_proj.dtype == torch.float8_e4m3fn
+        compute_dtype = x.dtype  # bf16 in normal use
+
+        outs = []
+        for i, xe in enumerate(x_per_expert):
+            if is_fp8:
+                w_gu = self.gate_up_proj[i].to(compute_dtype) * self.gate_up_proj_scale[i].to(compute_dtype)
+                w_dn = self.down_proj[i].to(compute_dtype) * self.down_proj_scale[i].to(compute_dtype)
+            else:
+                w_gu = self.gate_up_proj[i]
+                w_dn = self.down_proj[i]
+            gate_up = torch.matmul(xe, w_gu)
+            gate, up = gate_up.chunk(2, dim=-1)
+            outs.append(torch.matmul(F.silu(gate) * up, w_dn))
+
+        out = torch.cat(outs, dim=0)
+        return torch.vstack((out, out.new_zeros((n_pad, out.shape[-1]))))
+
+    def patched_forward(self, x, num_tokens_per_expert):
+        return patched_for_loop(self, x, num_tokens_per_expert)
+
+    cls.__init__ = patched_init
+    cls._run_experts_for_loop = patched_for_loop
+    cls.forward = patched_forward
+    setattr(cls, _PATCH_FLAG, True)
+
+
+def _patched_fp8_linear_forward(self, x):
+    # Dequantize qdata fp8 + per-output-channel scale on the fly, then F.linear.
+    w = self.weight.to(x.dtype) * self._fp8_scale.to(x.dtype)
+    return F.linear(x, w, self.bias)
+
+
+def install_fp8_linear(linear: nn.Linear, fp8_qdata: torch.Tensor, scale: torch.Tensor) -> None:
+    """Convert one Linear instance to fp8-stored, per-instance forward override."""
+    assert isinstance(linear, nn.Linear), f"expected nn.Linear, got {type(linear).__name__}"
+    linear.weight = nn.Parameter(fp8_qdata, requires_grad=False)
+    linear.register_buffer("_fp8_scale", scale.to(SCALE_DTYPE))
+    # Bind forward at the instance level — does NOT affect other nn.Linear instances.
+    linear.forward = _patched_fp8_linear_forward.__get__(linear, type(linear))
+
+
+def load_fp8_safetensors_transformer(safetensors_path: str, config_dir: str | None = None):
+    """
+    Load a Nucleus-Image FP8 transformer from a single safetensors file. The safetensors
+    holds raw fp8 + per-channel bf16 scale tensors (no torchao wrapper objects).
+
+    SwiGLUExperts: weights `gate_up_proj` / `down_proj` (fp8) + buffers `*_scale` (bf16).
+    Standard nn.Linear: weight (fp8) + buffer `_fp8_scale` (bf16) — forward override
+    installed by `install_fp8_linear` so the dequantize+matmul happens inline.
+    """
+    from pathlib import Path
+    import json
+    from safetensors.torch import safe_open
+    from diffusers import AutoModel
+    from accelerate import init_empty_weights
+
+    apply_patch()
+
+    cfg_dir = Path(config_dir) if config_dir else Path(safetensors_path).parent
+    cfg = json.loads((cfg_dir / "config.json").read_text(encoding="utf-8"))
+
+    # Locate the actual transformer class via the AutoModel mapping.
+    from diffusers import NucleusMoEImageTransformer2DModel
+    with init_empty_weights():
+        model = NucleusMoEImageTransformer2DModel.from_config(cfg)
+
+    fp8_dtypes = (torch.float8_e4m3fn, torch.float8_e5m2)
+
+    # Two passes over the file:
+    # Pass 1: identify Linears that have an `_fp8_scale` sibling — install fp8 weight + scale + forward.
+    # Pass 2: load everything else via standard assign.
+    with safe_open(safetensors_path, framework="pt", device="cpu") as f:
+        all_keys = set(f.keys())
+        # Linear keys with fp8 representation: those with a `_fp8_scale` buffer.
+        linear_paths = set()
+        for k in all_keys:
+            if k.endswith("._fp8_scale"):
+                base = k[: -len("._fp8_scale")]
+                if (base + ".weight") in all_keys:
+                    linear_paths.add(base)
+
+        # Pass 1: install fp8 Linears.
+        for base in linear_paths:
+            module = model.get_submodule(base)
+            qdata = f.get_tensor(base + ".weight").to("cpu")
+            scale = f.get_tensor(base + "._fp8_scale").to("cpu")
+            bias = None
+            if (base + ".bias") in all_keys:
+                bias = f.get_tensor(base + ".bias").to("cpu").to(torch.bfloat16)
+            # The init_empty_weights model has meta-device weights; replace cleanly.
+            if module.bias is not None and bias is not None:
+                module.bias = nn.Parameter(bias)
+            install_fp8_linear(module, qdata, scale)
+
+        # Pass 2: rest of the state dict (SwiGLUExperts fp8 + scales, norms, embeddings, etc.).
+        for k in all_keys:
+            # Skip keys we already handled in pass 1.
+            if any(k == p + ".weight" or k == p + "._fp8_scale" or k == p + ".bias" for p in linear_paths):
+                continue
+            t = f.get_tensor(k).to("cpu")
+            module_path, _, attr = k.rpartition(".")
+            module = model.get_submodule(module_path) if module_path else model
+            cur = getattr(module, attr, None)
+            if isinstance(cur, nn.Parameter) or cur is None:
+                # cur may be a meta-device Parameter — replace whole thing
+                if t.dtype.is_floating_point and t.dtype not in fp8_dtypes and t.dtype != torch.bfloat16:
+                    t = t.to(torch.bfloat16)
+                setattr(module, attr, nn.Parameter(t, requires_grad=False) if attr in dict(module.named_parameters(recurse=False)) or cur is None or isinstance(cur, nn.Parameter) else t)
+            else:
+                # Buffer
+                if t.dtype.is_floating_point and t.dtype not in fp8_dtypes and t.dtype != torch.bfloat16:
+                    t = t.to(torch.bfloat16)
+                module.register_buffer(attr, t)
+
+    # Final cleanup: any param/buffer still on meta should never happen, but assert.
+    meta_left = [n for n, p in model.named_parameters() if p.is_meta]
+    if meta_left:
+        raise RuntimeError(f"Some parameters still on meta device after load: {meta_left[:5]}... ({len(meta_left)} total)")
+
+    print(f"  load_fp8_safetensors_transformer: installed {len(linear_paths)} fp8 Linears")
+    return model
+
+
+def load_fp8_transformer(model_dir: str):
+    """
+    Load an FP8-quantized Nucleus transformer from `model_dir` while preserving on-disk
+    fp8_e4m3fn dtypes for SwiGLUExperts weights. Other floating params are normalized to bf16.
+
+    Why: AutoModel.from_pretrained(torch_dtype=torch.bfloat16) force-casts ALL floating
+    weights including fp8. We do the standard load (which casts to bf16), then re-stream
+    the on-disk shards and reassign every fp8_e4m3fn tensor as a fresh nn.Parameter so
+    the dtype is preserved. TorchAO Float8Tensor wrappers for nn.Linear are restored
+    correctly by the standard loader (its DiffusersAutoQuantizer hook intercepts them).
+    """
+    import json
+    from pathlib import Path
+    from diffusers import AutoModel
+
+    apply_patch()  # ensure SwiGLUExperts is patched before construction
+
+    model = AutoModel.from_pretrained(
+        model_dir,
+        use_safetensors=False,
+        low_cpu_mem_usage=True,
+    )
+
+    # Re-stream disk to recover fp8 dtypes (auto-cast lost them).
+    md = Path(model_dir)
+    idx_path = md / "diffusion_pytorch_model.bin.index.json"
+    if idx_path.exists():
+        idx = json.loads(idx_path.read_text(encoding="utf-8"))
+        files = sorted(set(idx["weight_map"].values()))
+    else:
+        files = ["diffusion_pytorch_model.bin"]
+
+    fp8_dtypes = (torch.float8_e4m3fn, torch.float8_e5m2)
+    fp8_reassigned = 0
+    for fname in files:
+        shard = torch.load(md / fname, map_location="cpu", weights_only=False)
+        for key, tensor in shard.items():
+            if not hasattr(tensor, "dtype") or tensor.dtype not in fp8_dtypes:
+                continue
+            module_path, _, attr = key.rpartition(".")
+            module = model.get_submodule(module_path)
+            cur = getattr(module, attr, None)
+            if isinstance(cur, nn.Parameter):
+                setattr(module, attr, nn.Parameter(tensor, requires_grad=False))
+            else:
+                # buffer
+                module.register_buffer(attr, tensor)
+            fp8_reassigned += 1
+        del shard
+
+    # Final pass: any non-fp8 floating param in fp32/fp16 → bf16 (uniformity).
+    for _, p in model.named_parameters():
+        if p.dtype in fp8_dtypes:
+            continue
+        if p.dtype.is_floating_point and p.dtype != torch.bfloat16:
+            p.data = p.data.to(torch.bfloat16)
+    for _, b in model.named_buffers():
+        if b.dtype in fp8_dtypes:
+            continue
+        if b.dtype.is_floating_point and b.dtype != torch.bfloat16:
+            b.data = b.data.to(torch.bfloat16)
+
+    print(f"  load_fp8_transformer: re-assigned {fp8_reassigned} fp8 tensors from shards")
+    return model
+
+
+@torch.no_grad()
+def quantize_swiglu_experts_(module: nn.Module) -> dict:
+    """Quantize a single SwiGLUExperts module in-place. Returns a small report dict."""
+    assert type(module).__name__ == "SwiGLUExperts", f"expected SwiGLUExperts, got {type(module).__name__}"
+    device = module.gate_up_proj.device
+
+    def _quant(w: torch.Tensor) -> tuple[torch.Tensor, torch.Tensor]:
+        # w shape: (num_experts, in_dim, out_dim). Per-expert per-output-channel scale on dim=-2.
+        w32 = w.detach().to(torch.float32)
+        scale = w32.abs().amax(dim=-2, keepdim=True).clamp(min=1e-12) / FP8_E4M3_MAX
+        q = (w32 / scale).clamp(-FP8_E4M3_MAX, FP8_E4M3_MAX).to(torch.float8_e4m3fn)
+        return q, scale.to(SCALE_DTYPE)
+
+    q_gu, s_gu = _quant(module.gate_up_proj.data)
+    q_dn, s_dn = _quant(module.down_proj.data)
+
+    # Reconstruction error gate. Use rel L2 (||W_hat - W||_2 / ||W||_2) per tensor — the right
+    # metric for forward fidelity (matmul amplifies L2-style error, not max-abs).
+    w_gu_orig = module.gate_up_proj.data.to(torch.float32)
+    w_dn_orig = module.down_proj.data.to(torch.float32)
+    rec_gu = q_gu.to(torch.float32) * s_gu.to(torch.float32)
+    rec_dn = q_dn.to(torch.float32) * s_dn.to(torch.float32)
+
+    rep = {
+        "gu_rel_l2": ((rec_gu - w_gu_orig).norm() / w_gu_orig.norm().clamp(min=1e-12)).item(),
+        "dn_rel_l2": ((rec_dn - w_dn_orig).norm() / w_dn_orig.norm().clamp(min=1e-12)).item(),
+        # Keep the loose rel-max-err for visibility but no longer used as a gate.
+        "gu_rel_max": ((rec_gu - w_gu_orig).abs().amax() / w_gu_orig.abs().amax().clamp(min=1e-12)).item(),
+        "dn_rel_max": ((rec_dn - w_dn_orig).abs().amax() / w_dn_orig.abs().amax().clamp(min=1e-12)).item(),
+    }
+
+    # In-place replacement
+    module.gate_up_proj = nn.Parameter(q_gu.to(device), requires_grad=False)
+    module.down_proj = nn.Parameter(q_dn.to(device), requires_grad=False)
+    module.gate_up_proj_scale.data = s_gu.to(device)
+    module.down_proj_scale.data = s_dn.to(device)
+    return rep