# TinyFlux Inference - Colab Cell
# 
# Install and run:
#   !pip install git+https://github.com/AbstractPhil/tinyflux.git
#   # Then run this cell

#@title TinyFlux Inference
prompt = "object, a red apple" #@param {type:"string"}
negative_prompt = "glass, cup" #@param {type:"string"}
seed = 42 #@param {type:"integer"}
steps = 50 #@param {type:"slider", min:4, max:50, step:1}
guidance = 5 #@param {type:"slider", min:1.0, max:15.0, step:0.5}

import torch
from tinyflux.model.zoo import ModelZoo
from tinyflux.trainer.sampling import Sampler
from PIL import Image
from IPython.display import display

device = "cuda" if torch.cuda.is_available() else "cpu"
dtype = torch.bfloat16 if device == "cuda" else torch.float32

# Load models
print("Loading models...")
zoo = ModelZoo(device=device, dtype=dtype)
zoo.load_vae()
zoo.load_clip()
zoo.load_t5()
model = zoo.load_tinyflux(
    source="AbstractPhil/tiny-flux-deep",
    load_ema=False, # this enables loading the checkpoint from the ema path, leave false for now to play with the most recent good version
    ema_path="checkpoints/step_409244.safetensors",
)
print("✓ Models loaded")

# Create sampler
sampler = Sampler(
    zoo=zoo,
    model=model,
    num_steps=steps,
    guidance_scale=guidance,
    shift=3.0,
    device=device,
    dtype=dtype,
)

# Generate
print(f"\nGenerating: {prompt}")
images = sampler.generate([prompt], seed=seed, negative_prompt=negative_prompt)

# Display
img = (images[0].permute(1, 2, 0).cpu().float().numpy() * 255).astype("uint8")
display(Image.fromarray(img))
print(f"Seed: {seed}, Steps: {steps}, CFG: {guidance}")