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

# Define available models and their corresponding Hugging Face repositories
MODEL_REPOS = {
    "Stable Diffusion XL Base 1.0": "stabilityai/stable-diffusion-xl-base-1.0",
    "SDXL-Turbo": "stabilityai/sdxl-turbo",
    "Playground v2 1024px Aesthetic": "playgroundai/playground-v2-1024px-aesthetic",
    "Segmind Vega": "segmind/Segmind-Vega",
    "SSD-1B": "segmind/SSD-1B",
    "Kandinsky 3": "kandinsky-community/kandinsky-3",
    "PixArt-LCM-XL-2-1024-MS": "PixArt-alpha/PixArt-LCM-XL-2-1024-MS",
    "BLIP Diffusion": "salesforce/blipdiffusion",
    "Muse-512-Finetuned": "amused/muse-512-finetuned",
    "Flux 1 Dev": "black-forest-labs/FLUX.1-dev"
}

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

# Cache for loaded pipelines
loaded_pipelines = {}

# Maximum seed value
MAX_SEED = np.iinfo(np.int32).max

def load_pipeline(model_name):
    """Load and cache the pipeline for the selected model."""
    if model_name in loaded_pipelines:
        return loaded_pipelines[model_name]
    repo_id = MODEL_REPOS[model_name]
    try:
        pipeline = DiffusionPipeline.from_pretrained(repo_id, torch_dtype=torch_dtype)
        pipeline.to(device)
        loaded_pipelines[model_name] = pipeline
        return pipeline
    except Exception as e:
        raise RuntimeError(f"Failed to load model '{model_name}': {e}")

def generate_image(prompt, model_name, width, height, guidance_scale, num_inference_steps, seed, randomize_seed):
    """Generate an image using the selected model and parameters."""
    if randomize_seed:
        seed = random.randint(0, MAX_SEED)
    generator = torch.Generator(device=device).manual_seed(seed)
    pipeline = load_pipeline(model_name)
    try:
        image = pipeline(
            prompt=prompt,
            guidance_scale=guidance_scale,
            num_inference_steps=num_inference_steps,
            width=width,
            height=height,
            generator=generator
        ).images[0]
        return image, seed
    except Exception as e:
        raise RuntimeError(f"Image generation failed: {e}")

# Define the Gradio interface
with gr.Blocks() as demo:
    gr.Markdown("# 🖼️ Text-to-Image Generator with Multiple Models")
    with gr.Row():
        with gr.Column():
            prompt = gr.Textbox(label="Prompt", placeholder="Enter your prompt here")
            model_name = gr.Dropdown(label="Select Model", choices=list(MODEL_REPOS.keys()), value="Stable Diffusion XL Base 1.0")
            width = gr.Slider(label="Width", minimum=256, maximum=1024, step=64, value=512)
            height = gr.Slider(label="Height", minimum=256, maximum=1024, step=64, value=512)
            guidance_scale = gr.Slider(label="Guidance Scale", minimum=1.0, maximum=20.0, step=0.5, value=7.5)
            num_inference_steps = gr.Slider(label="Inference Steps", minimum=1, maximum=100, step=1, value=50)
            seed = gr.Slider(label="Seed", minimum=0, maximum=MAX_SEED, step=1, value=0)
            randomize_seed = gr.Checkbox(label="Randomize Seed", value=True)
            generate_button = gr.Button("Generate Image")
        with gr.Column():
            output_image = gr.Image(label="Generated Image")
            output_seed = gr.Textbox(label="Used Seed", interactive=False)

    generate_button.click(
        fn=generate_image,
        inputs=[prompt, model_name, width, height, guidance_scale, num_inference_steps, seed, randomize_seed],
        outputs=[output_image, output_seed]
    )

if __name__ == "__main__":
    demo.launch()