Upload generate.py with huggingface_hub

generate.py

@@ -0,0 +1,108 @@
+#!/usr/bin/env python3
+"""Generate images with Nucleus-Image on Apple Silicon (MLX).
+
+Usage:
+    python generate.py --prompt "A red apple on a white table"
+    python generate.py --prompt "A futuristic city at sunset" --steps 30 --seed 42 --output city.png
+"""
+
+import argparse
+import gc
+import time
+
+import mlx.core as mx
+import torch
+from transformers import AutoModel, AutoProcessor
+
+from nucleus_image.pipeline import NucleusImagePipeline
+
+SYSTEM_PROMPT = "You are an image generation assistant."
+TEXT_MODEL_ID = "NucleusAI/Nucleus-Image"
+HIDDEN_LAYER_INDEX = -8  # 8th from last hidden state
+
+
+def encode_text(prompt: str, processor, text_model) -> mx.array:
+    """Encode a text prompt into embeddings using the chat template format."""
+    messages = [
+        {"role": "system", "content": SYSTEM_PROMPT},
+        {"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[HIDDEN_LAYER_INDEX][0]
+    return mx.array(hidden.cpu().float().numpy())
+
+
+def main():
+    parser = argparse.ArgumentParser(description="Generate images with MLX Nucleus-Image")
+    parser.add_argument("--prompt", type=str, required=True, help="Text prompt for image generation")
+    parser.add_argument("--height", type=int, default=512, help="Image height (default: 512)")
+    parser.add_argument("--width", type=int, default=512, help="Image width (default: 512)")
+    parser.add_argument("--steps", type=int, default=50, help="Number of inference steps (default: 50)")
+    parser.add_argument("--cfg", type=float, default=4.0, help="Classifier-free guidance scale (default: 4.0)")
+    parser.add_argument("--seed", type=int, default=None, help="Random seed for reproducibility")
+    parser.add_argument("--output", type=str, default="output.png", help="Output file path (default: output.png)")
+    parser.add_argument("--quantize", type=int, default=4, choices=[4, 8, None], help="DiT quantization bits (default: 4)")
+    args = parser.parse_args()
+
+    t_total = time.time()
+
+    # Step 1: Load text encoder and encode prompt + negative (empty string)
+    print("Loading text encoder...")
+    t0 = time.time()
+    processor = AutoProcessor.from_pretrained(
+        TEXT_MODEL_ID, subfolder="processor", trust_remote_code=True
+    )
+    text_model = AutoModel.from_pretrained(
+        TEXT_MODEL_ID, subfolder="text_encoder",
+        dtype=torch.bfloat16, trust_remote_code=True,
+    )
+    text_model.eval()
+    print(f"  Text encoder loaded in {time.time() - t0:.1f}s")
+
+    print("Encoding prompt...")
+    t0 = time.time()
+    text_emb = encode_text(args.prompt, processor, text_model)
+
+    print("Encoding negative embeddings...")
+    neg_emb = encode_text("", processor, text_model)
+    print(f"  Text encoding done in {time.time() - t0:.1f}s")
+
+    # Free text encoder memory (~16GB)
+    del text_model, processor
+    gc.collect()
+
+    # Step 2: Load MLX pipeline (DiT + VAE)
+    print(f"Loading MLX pipeline (quantize={args.quantize})...")
+    t0 = time.time()
+    pipe = NucleusImagePipeline.from_pretrained(quantize=args.quantize)
+    print(f"  Pipeline loaded in {time.time() - t0:.1f}s")
+
+    # Step 3: Generate image
+    print(f"Generating {args.height}x{args.width}, {args.steps} steps, CFG {args.cfg}...")
+    img = pipe.generate(
+        text_embeddings=mx.expand_dims(text_emb, 0),
+        neg_text_embeddings=mx.expand_dims(neg_emb, 0),
+        height=args.height,
+        width=args.width,
+        num_inference_steps=args.steps,
+        guidance_scale=args.cfg,
+        seed=args.seed,
+    )
+
+    img.save(args.output)
+    print(f"Saved to {args.output}")
+    print(f"Total time: {time.time() - t_total:.1f}s")
+
+
+if __name__ == "__main__":
+    main()