import gradio as gr
import numpy as np
import random
import torch
import os
from diffusers import DiffusionPipeline

# --------------------------
# 0. Read HF Token
# --------------------------
HF_TOKEN = os.environ.get("HF_TOKEN")
if HF_TOKEN is None:
    raise ValueError("❌ ERROR: HF_TOKEN is not set as a secret in your Space!")

# --------------------------
# 1. Base model + LoRA
# --------------------------
BASE_MODEL = "black-forest-labs/FLUX.1-dev"
LORA_MODEL = "lamanabin/cinderella-flux-lora"   # <-- FIXED REPO ID

device = "cuda" if torch.cuda.is_available() else "cpu"
torch_dtype = torch.float16 if device == "cuda" else torch.float32


# --------------------------
# 2. Load FLUX.1-dev
# --------------------------
print("🚀 Loading FLUX.1-dev...")
pipe = DiffusionPipeline.from_pretrained(
    BASE_MODEL,
    use_auth_token=HF_TOKEN,
    torch_dtype=torch_dtype,
    use_safetensors=True,
).to(device)


# --------------------------
# 3. Load your LoRA
# --------------------------
print("🎨 Loading your LoRA:", LORA_MODEL)

pipe.load_lora_weights(
    LORA_MODEL,
    use_auth_token=HF_TOKEN,
)

pipe.fuse_lora()
print("✅ FLUX + LoRA loaded successfully!")


MAX_SEED = np.iinfo(np.int32).max


# --------------------------
# 4. Inference function
# --------------------------
def infer(prompt, seed, randomize_seed, num_inference_steps):
    if randomize_seed:
        seed = random.randint(0, MAX_SEED)

    generator = torch.Generator(device=device).manual_seed(seed)

    image = pipe(
        prompt=prompt,
        guidance_scale=0.0,
        num_inference_steps=num_inference_steps,
        generator=generator,
    ).images[0]

    return image, seed


# --------------------------
# 5. Gradio UI
# --------------------------
with gr.Blocks() as demo:
    gr.Markdown("# 🎨 FLUX.1-dev + Cinderella LoRA")

    prompt = gr.Textbox(label="Prompt", lines=2)
    run_button = gr.Button("Generate")
    result = gr.Image(type="pil")

    seed = gr.Slider(0, MAX_SEED, value=0, label="Seed")
    randomize_seed = gr.Checkbox(label="Random Seed", value=True)
    num_steps = gr.Slider(1, 48, value=28, step=1, label="Inference Steps")

    run_button.click(
        fn=infer,
        inputs=[prompt, seed, randomize_seed, num_steps],
        outputs=[result, seed]
    )

demo.launch()