"""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}")