import gradio as gr
import numpy as np
import random
from diffusers import DiffusionPipeline, DDIMScheduler, LMSDiscreteScheduler
import torch

# Device configuration (explicitly set to CPU)
DEVICE = "cpu"

# Model Options (for user selection)
MODEL_OPTIONS = {
    "High Quality (Slower)": "stabilityai/stable-diffusion-xl-base-1.0",
    "Fast (Lower Quality)": "CompVis/stable-diffusion-v1-4",  # Smaller, faster model
}

# Default to faster model
DEFAULT_MODEL_ID = MODEL_OPTIONS["Fast (Lower Quality)"]

def generate_image(prompt, negative_prompt, seed, randomize_seed, width, height, guidance_scale, num_inference_steps, num_images, model_choice):
    model_id = MODEL_OPTIONS[model_choice]

    # Load Model based on user selection
    pipe = DiffusionPipeline.from_pretrained(
        model_id, torch_dtype=torch.float32
    )

    # Use LMSDiscreteScheduler for faster generation on CPU with CompVis model
    if model_choice == "Fast (Lower Quality)":
        pipe.scheduler = LMSDiscreteScheduler.from_config(pipe.scheduler.config)

    pipe = pipe.to(DEVICE)

    generator = torch.Generator(device=DEVICE)
    if not randomize_seed:
        generator = generator.manual_seed(seed)

    images = pipe(
        prompt,
        negative_prompt=negative_prompt,
        width=width,
        height=height,
        guidance_scale=guidance_scale,
        num_inference_steps=num_inference_steps,
        num_images_per_prompt=num_images,
        generator=generator,
    ).images

    return images

# Gradio Interface
with gr.Blocks(title="Генерація зображень за текстом", theme=gr.themes.Soft()) as demo:
    gr.Markdown(
        """
        ## Text-to-Image Generation 🤖🎨
        **Створіть дивовижні зображення зі своєї уяви!**  
        Введіть опис, налаштуйте параметри і дозвольте моделі створити для вас витвір мистецтва.
        """)

    with gr.Row():
        prompt = gr.Textbox(label="Опис", placeholder="Напишіть ваш опис тут...")
        negative_prompt = gr.Textbox(label="Негативний опис (необов'язково)")

    with gr.Row():
        seed = gr.Number(label="Початкове число", value=0)
        randomize_seed = gr.Checkbox(label="Випадкове початкове число", value=True)
        
    with gr.Row():  # Added this row for model selection
        model_choice = gr.Radio(label="Виберіть модель", choices=list(MODEL_OPTIONS.keys()), value=DEFAULT_MODEL_ID) 

    with gr.Row():
        width = gr.Slider(label="Ширина", minimum=256, maximum=MAX_IMAGE_SIZE, value=512, step=64)
        height = gr.Slider(label="Висота", minimum=256, maximum=MAX_IMAGE_SIZE, value=512, step=64)
  
    with gr.Accordion("Додаткові налаштування", open=False):
        with gr.Row():
            guidance_scale = gr.Slider(label="Рівень відповідності опису", minimum=0.0, maximum=20.0, value=7.5, step=0.1, info="Наскільки точно модель повинна слідувати опису.")
            num_inference_steps = gr.Slider(label="Кількість кроків", minimum=10, maximum=100, value=50, step=5, info="Більше кроків може покращити якість, але займе більше часу.")
            num_images = gr.Slider(label="Кількість зображень", minimum=1, maximum=4, value=1, step=1)

    run_button = gr.Button("Згенерувати")
    gallery = gr.Gallery(label="Згенеровані зображення")

    run_button.click(
        fn=generate_image,
        inputs=[prompt, negative_prompt, seed, randomize_seed, width, height, guidance_scale, num_inference_steps, num_images, model_choice],
        outputs=gallery,
    )

demo.launch(debug=True)