import gradio as gr
from diffusers import StableDiffusionPipeline
import torch

# Словарь с путями к моделям
model_paths = {
    "Vector Art": "stablediffusionapi/vector-art",
    # Добавьте другие модели здесь
}

# Инициализация и загрузка моделей
models = {}
for name, path in model_paths.items():
    pipeline = StableDiffusionPipeline.from_pretrained(path)
    pipeline.to("cpu")
    models[name] = pipeline

# Функция для генерации изображения
def generate_image(prompt, model_name, sampling_steps, cfg, sampling_method, seed, width, height):
    # Выбор модели из словаря
    model = models[model_name]
    # Установка параметров генерации
    generator = torch.manual_seed(seed) if seed != -1 else None
    # Генерация изображения
    image = model(prompt, num_inference_steps=sampling_steps, guidance_scale=cfg, generator=generator).images[0]
    # Преобразование изображения для отображения в Gradio
    return image

# Создание интерфейса с использованием gr.Blocks
with gr.Blocks() as app:
    gr.Markdown("Генератор изображений")
    with gr.Tab("Основные настройки"):
        with gr.Row():
            prompt_input = gr.Textbox(label="Введите текстовый запрос")
            model_select = gr.Radio(choices=list(model_paths.keys()), label="Выберите модель", value="Vector Art")
    with gr.Tab("Расширенные настройки"):
        with gr.Row():
            sampling_steps_input = gr.Slider(minimum=1, maximum=30, step=1, label="Sampling Steps", value=20)
            cfg_input = gr.Slider(minimum=1, maximum=20, step=1, label="Classifier Free Guidance (CFG)", value=7)
            sampling_method_input = gr.Radio(choices=["DDIM", "DDPM", "PLMS"], label="Sampling Method", value="DDIM")
            seed_input = gr.Slider(minimum=-1, maximum=100000, step=1, label="Seed", value=-1)
            width_input = gr.Slider(minimum=15, maximum=1023, step=1, label="Width", value=512)
            height_input = gr.Slider(minimum=15, maximum=1023, step=1, label="Height", value=512)
    with gr.Row():
        generate_button = gr.Button("Сгенерировать")
    output_image = gr.Image(label="Сгенерированное изображение")

    # Подключение функции генерации к кнопке
    generate_button.click(
        generate_image,
        inputs=[prompt_input, model_select, sampling_steps_input, cfg_input, sampling_method_input, seed_input, width_input, height_input],
        outputs=output_image
    )

app.launch()