import gradio as gr
from diffusers import StableDiffusionPipeline, DPMSolverMultistepScheduler
import torch

# --- Configuration ---
HF_REPO_ID = "aanchal77/Final-One"
BASE_MODEL_ID = "runwayml/stable-diffusion-v1-5"

# --- Define Available LoRAs ---
AVAILABLE_LORAS = {
    "None (Base Model)": None,
    # --- Artists ---
    "Artist: Vincent van Gogh": "artists/Vincent_van_Gogh",
    "Artist: Claude Monet": "artists/Claude_Monet",
    "Artist: Rembrandt": "artists/Rembrandt",
    "Artist: Pablo Picasso": "artists/Pablo_Picasso",
    # --- Styles ---
    "Style: Impressionism": "styles/Impressionism",
    "Style: Baroque": "styles/Baroque",
    "Style: Cubism": "styles/Cubism",
    "Style: Abstract Expressionism": "styles/Abstract_Expressionism",
    "Style: Romanticism": "styles/Romanticism",
    "Style: Realism": "styles/Realism",
    "Style: Post Impressionism": "styles/Post_Impressionism",
}
print("✅ LoRA models from your Hugging Face repo are configured.")

# --- Setup ---
device = "cuda" if torch.cuda.is_available() else "cpu"
dtype = torch.float16 if device == "cuda" else torch.float32
print(f"Using device: {device}")

print(f"🎨 Loading base model: {BASE_MODEL_ID}")
pipe = StableDiffusionPipeline.from_pretrained(BASE_MODEL_ID, torch_dtype=dtype)

# 🔧 Replace the fragile PNDM scheduler with a robust one to avoid index/NoneType errors
pipe.scheduler = DPMSolverMultistepScheduler.from_config(pipe.scheduler.config)

pipe = pipe.to(device)
if device == "cpu":
    pipe.enable_attention_slicing()

# --- The Core Generation Function ---
def generate(prompt, quality, lora_choice):
    # Reset to base weights
    pipe.unload_lora_weights()

    lora_subfolder = AVAILABLE_LORAS.get(lora_choice)

    if lora_subfolder:
        print(f"✨ Downloading and applying LoRA: {lora_choice}")
        try:
            pipe.load_lora_weights(
                HF_REPO_ID,
                subfolder=lora_subfolder,
                weight_name="adapter_model.safetensors"  # ensure exact file
            )
        except Exception as e:
            print(f"❌ Failed to load LoRA from Hub '{HF_REPO_ID}/{lora_subfolder}': {e}")
    else:
        print("🎨 Using base model (no LoRA selected)")

    steps = 25 if quality == "Fast" else 40
    guidance_scale = 7.5

    print(f"🚀 Generating with prompt: '{prompt}'")
    with torch.no_grad():
        image = pipe(
            prompt,
            num_inference_steps=steps,
            guidance_scale=guidance_scale
        ).images[0]

    return image

# --- Build the Gradio UI ---
title = f"🎨 Stable Diffusion Gallery from {HF_REPO_ID}"
description = (
    "Select a trained LoRA model from your Hugging Face repository to apply its style. "
    "The first time you select a LoRA, it may take a moment to download."
)

demo = gr.Interface(
    fn=generate,
    inputs=[
        gr.Textbox(label="Enter your prompt", placeholder="A beautiful painting of a fantasy landscape..."),
        gr.Dropdown(["Fast", "High Quality"], value="Fast", label="Generation Quality"),
        gr.Dropdown(
            choices=list(AVAILABLE_LORAS.keys()),
            value="None (Base Model)",
            label="Select a Trained LoRA Model"
        )
    ],
    outputs=gr.Image(label="Generated Image"),
    title=title,
    description=description,
    examples=[
        ["A portrait of an astronaut, cinematic lighting, by vincent van gogh", "Fast", "Artist: Vincent van Gogh"],
        ["A peaceful village in the mountains, impressionism style", "High Quality", "Style: Impressionism"],
    ],
    cache_examples=False,   # 🔒 prevent startup 500s if an example errors
)

demo.launch()