Upload test_tiny.py with huggingface_hub

test_tiny.py

@@ -0,0 +1,43 @@
+"""Tiny test: 128x128, 4 steps, no CFG, with text embeddings."""
+import sys, time, torch, numpy as np
+import mlx.core as mx
+sys.path.insert(0, ".")
+
+# Extract text embeddings
+from transformers import AutoProcessor, AutoModel
+PROMPT = "A red apple on a white table"
+SYSTEM = "You are an image generation assistant."
+
+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()
+
+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]
+text_emb = mx.array(hidden.cpu().float().numpy())
+print(f"Text: {text_emb.shape}")
+
+del text_model, processor
+import gc; gc.collect()
+
+# Generate
+from nucleus_image.pipeline import NucleusImagePipeline
+pipe = NucleusImagePipeline.from_pretrained(quantize=4)
+
+print("Generating 128x128, 4 steps, no CFG...")
+t0 = time.time()
+img = pipe.generate(
+    text_embeddings=mx.expand_dims(text_emb, 0),
+    height=128, width=128,
+    num_inference_steps=4,
+    guidance_scale=1.0,  # no CFG for speed
+    seed=42,
+)
+print(f"Done in {time.time()-t0:.1f}s")
+img.save("/Users/ritesh/Dev/model-training/nucleus-image/mlx/test-output/tiny_test.png")
+print(f"Saved! {img.size}")