Upload README.md with huggingface_hub

README.md

@@ -17,35 +17,42 @@ pipeline_tag: text-to-image
 
 An [MLX](https://github.com/ml-explore/mlx) port of [NucleusAI/Nucleus-Image](https://huggingface.co/NucleusAI/Nucleus-Image), a 17B parameter Mixture-of-Experts (MoE) DiT for text-to-image generation. Runs natively on Apple Silicon.
 
-## Sample Outputs (512x512, 50 steps, CFG 4.0, 4-bit)
+## Sample Outputs (512x512, 30 steps, CFG 4.0, 4-bit)
 
-*Run `generate.py` to produce your own samples in the `samples/` directory.*
+| "A red apple on a white table" | "A golden retriever puppy playing in autumn leaves" | "A futuristic city skyline at sunset" |
+|:---:|:---:|:---:|
+| <img src="samples/apple.png" width="256"> | <img src="samples/puppy.png" width="256"> | <img src="samples/city.png" width="256"> |
 
-## Architecture
-
-| Component | Details |
-|-----------|---------|
-| DiT | 17B total params, ~2B active per token |
-| Layers | 32 (3 dense + 29 MoE) |
-| Experts | 64 routed + 1 shared per MoE layer |
-| Routing | Expert-choice (capacity-based) |
-| Attention | GQA: 16 query / 4 KV heads, head_dim=128 |
-| Text Encoder | Qwen3-VL-8B-Instruct (PyTorch, hybrid) |
-| VAE | AutoencoderKLQwenImage, 16-ch latents |
-| Scheduler | Flow Matching Euler |
+| "A steaming cup of coffee on a rainy windowsill" | "An astronaut riding a horse on the moon" |
+|:---:|:---:|
+| <img src="samples/coffee.png" width="256"> | <img src="samples/astronaut.png" width="256"> |
 
 ## Quick Start
 
+### 1. Clone the repo
+
+```bash
+git clone https://huggingface.co/treadon/mlx-nucleus-image
+cd mlx-nucleus-image
+```
+
+### 2. Install dependencies
+
 ```bash
 pip install mlx torch transformers huggingface_hub pillow
 ```
 
-### CLI (recommended)
+### 3. Generate an image
 
 ```bash
 python generate.py --prompt "A red apple on a white table" --seed 42
+```
+
+The first run downloads ~34GB of weights (cached for subsequent runs).
+
+### More options
 
-# More options
+```bash
 python generate.py \
   --prompt "A futuristic city skyline at sunset" \
   --height 512 --width 512 \
@@ -54,56 +61,29 @@ python generate.py \
   --quantize 4
 ```
 
-### Python API
+| Flag | Default | Description |
+|------|---------|-------------|
+| `--prompt` | required | Text prompt |
+| `--height` | 512 | Image height |
+| `--width` | 512 | Image width |
+| `--steps` | 50 | Denoising steps |
+| `--cfg` | 4.0 | Guidance scale |
+| `--seed` | random | Random seed |
+| `--output` | output.png | Output file |
+| `--quantize` | 4 | Quantization bits (4, 8, or None) |
 
-```python
-import torch
-import mlx.core as mx
-from transformers import AutoProcessor, AutoModel
-from nucleus_image import NucleusImagePipeline
-
-# Step 1: Encode text (PyTorch — runs once, then freed)
-processor = AutoProcessor.from_pretrained(
-    "NucleusAI/Nucleus-Image", subfolder="processor", trust_remote_code=True
-)
-text_model = AutoModel.from_pretrained(
-    "NucleusAI/Nucleus-Image", subfolder="text_encoder",
-    dtype=torch.bfloat16, trust_remote_code=True
-)
-text_model.eval()
-
-PROMPT = "A red apple on a white table"
-SYSTEM = "You are an image generation assistant."
-messages = [
-    {"role": "system", "content": SYSTEM},
-    {"role": "user", "content": [{"type": "text", "text": PROMPT}]},
-]
-formatted = processor.apply_chat_template(messages, tokenize=False, add_generation_prompt=True)
-inputs = processor(text=[formatted], return_tensors="pt", padding=True)
-
-with torch.no_grad():
-    outputs = text_model(
-        input_ids=inputs["input_ids"],
-        attention_mask=inputs.get("attention_mask"),
-        output_hidden_states=True, use_cache=False,
-    )
-    hidden = outputs.hidden_states[-8][0]  # 8th from last
-text_emb = mx.array(hidden.cpu().float().numpy())
-
-del text_model, processor  # Free ~16GB
-
-# Step 2: Generate image (MLX)
-pipe = NucleusImagePipeline.from_pretrained(quantize=4)  # 4-bit DiT
+## Architecture
 
-img = pipe.generate(
-    text_embeddings=mx.expand_dims(text_emb, 0),
-    height=512, width=512,
-    num_inference_steps=30,
-    guidance_scale=4.0,
-    seed=42,
-)
-img.save("output.png")
-```
+| Component | Details |
+|-----------|---------|
+| DiT | 17B total params, ~2B active per token |
+| Layers | 32 (3 dense + 29 MoE) |
+| Experts | 64 routed + 1 shared per MoE layer |
+| Routing | Expert-choice (capacity-based) |
+| Attention | GQA: 16 query / 4 KV heads, head_dim=128 |
+| Text Encoder | Qwen3-VL-8B-Instruct (PyTorch, hybrid) |
+| VAE | AutoencoderKLQwenImage, 16-ch latents |
+| Scheduler | Flow Matching Euler |
 
 ## Performance (M4 Pro 64GB)
 
@@ -118,7 +98,7 @@ img.save("output.png")
 The port is a hybrid approach:
 - **Text encoder** stays in PyTorch (Qwen3-VL-8B, ~16GB). Loaded once to extract embeddings, then freed.
 - **DiT** (17B MoE) runs in MLX with optional 4-bit quantization for attention/modulation layers. Expert weights stay in bfloat16.
-- **VAE decoder** runs in MLX (254MB). Conv3d weights converted to Conv2d by extracting the last temporal kernel slice (CausalConv3d).
+- **VAE decoder** runs in MLX (~50MB converted). Original CausalConv3d weights converted to Conv2d.
 
 ### Key Conversion Details
 
@@ -131,21 +111,59 @@ The port is a hybrid approach:
 | AdaLayerNormContinuous | LayerNorm(affine=False) + adaptive scale/shift | scale first, shift second |
 | Timesteps(scale=1000) | `timestep_embedding(sigma * 1000)` | Pipeline normalizes t/1000 before DiT |
 
+## Python API
+
+```python
+import torch
+import mlx.core as mx
+from transformers import AutoProcessor, AutoModel
+from nucleus_image import NucleusImagePipeline
+
+# Encode text (PyTorch — loaded once, then freed)
+processor = AutoProcessor.from_pretrained("NucleusAI/Nucleus-Image", subfolder="processor", trust_remote_code=True)
+text_model = AutoModel.from_pretrained("NucleusAI/Nucleus-Image", subfolder="text_encoder", dtype=torch.bfloat16, trust_remote_code=True).eval()
+
+SYSTEM = "You are an image generation assistant."
+def encode(prompt):
+    messages = [{"role": "system", "content": SYSTEM}, {"role": "user", "content": [{"type": "text", "text": prompt}]}]
+    formatted = processor.apply_chat_template(messages, tokenize=False, add_generation_prompt=True)
+    inputs = processor(text=[formatted], return_tensors="pt", padding=True)
+    with torch.no_grad():
+        out = text_model(input_ids=inputs["input_ids"], attention_mask=inputs.get("attention_mask"),
+                        output_hidden_states=True, use_cache=False)
+    return mx.array(out.hidden_states[-8][0].cpu().float().numpy())
+
+text_emb = encode("A red apple on a white table")
+neg_emb = encode("")  # empty string for proper CFG
+del text_model, processor  # free ~16GB
+
+# Generate (MLX)
+pipe = NucleusImagePipeline.from_pretrained(quantize=4)
+img = pipe.generate(
+    text_embeddings=mx.expand_dims(text_emb, 0),
+    neg_text_embeddings=mx.expand_dims(neg_emb, 0),
+    height=512, width=512, num_inference_steps=30, guidance_scale=4.0, seed=42,
+)
+img.save("output.png")
+```
+
 ## Files
 
 ```
-generate.py           — CLI entry point (108 lines)
+generate.py              — CLI entry point
+samples/                 — Pre-generated sample images
 nucleus_image/
-  __init__.py         — Package exports
-  dit.py              — 17B MoE DiT (509 lines)
-  vae.py              — VAE decoder with Conv3d->Conv2d (189 lines)
-  pipeline.py         — End-to-end pipeline (184 lines)
-  scheduler.py        — Flow matching Euler scheduler (24 lines)
-  text_encoder.py     — Text encoder wrapper (optional)
+  dit.py                 — 17B MoE DiT
+  vae.py                 — VAE decoder (Conv3d→Conv2d)
+  pipeline.py            — End-to-end pipeline
+  scheduler.py           — Flow matching Euler scheduler
+dit/                     — Pre-converted DiT weights (safetensors)
+vae/                     — Pre-converted VAE weights (safetensors)
 ```
 
 ## Acknowledgments
 
 - [NucleusAI](https://huggingface.co/NucleusAI) for the original Nucleus-Image model
 - [Apple MLX](https://github.com/ml-explore/mlx) framework
-- Built by [@treadon](https://twitter.com/treadon)
+- Source code: [github.com/treadon/mlx-nucleus-image](https://github.com/treadon/mlx-nucleus-image)
+- Built by [@treadon](https://x.com/treadon)