import gradio as gr
from diffusers import StableDiffusionPipeline
import torch

# Список моделей для вибору
IMAGE_MODELS = {
    "Stable Diffusion v1.5": "runwayml/stable-diffusion-v1-5",
    "Stable Diffusion XL Base": "stabilityai/stable-diffusion-xl-base-1.0",
    "Kandinsky 2.2": "kandinsky-community/kandinsky-2-2-decoder",
    "DreamShaper XL": "Lykon/DreamShaper"
}

# Кеш для завантажених пайплайнів
loaded_pipes = {}

def get_pipe(model_name):
    if model_name not in loaded_pipes:
        pipe = StableDiffusionPipeline.from_pretrained(
            IMAGE_MODELS[model_name],
            torch_dtype=torch.float16 if torch.cuda.is_available() else torch.float32
        )
        pipe = pipe.to("cuda" if torch.cuda.is_available() else "cpu")
        loaded_pipes[model_name] = pipe
    return loaded_pipes[model_name]

def generate(prompt, model_name, steps, guidance):
    pipe = get_pipe(model_name)
    result = pipe(prompt, num_inference_steps=steps, guidance_scale=guidance)
    return result.images[0]

# Інтерфейс Gradio з вибором моделі та налаштуваннями якості
demo = gr.Interface(
    fn=generate,
    inputs=[
        gr.Textbox(label="Prompt", placeholder="Напиши опис сюди"),
        gr.Dropdown(choices=list(IMAGE_MODELS.keys()), value="Stable Diffusion XL Base", label="Модель"),
        gr.Slider(minimum=10, maximum=50, value=30, step=1, label="Кількість кроків (steps)"),
        gr.Slider(minimum=1.0, maximum=15.0, value=7.5, step=0.5, label="Guidance scale")
    ],
    outputs=gr.Image(label="Result"),
    title="AI Image Generator (Multi-Models + Quality Settings)"
)

demo.launch(ssr_mode=False)