app.py

import gradio as gr
import torch
from PIL import Image

# Список моделей
models = [
    "CLIP",
    "DALL-E 2",
    "VQGAN + CLIP",
    "VQGAN",
    "Dreams",
    "VQGAN-CLIP-ADA",
    "DALL-E Mini",
    "VQGAN-CLIP-Ada-Mini",
    "DALL-E Mini-Fusion",
    "VQGAN-CLIP-Ada-Mini-Fusion",
    "DALL-E 2 Mini",
    "VQGAN-CLIP-Ada-Mini-2",
]

# Функция генерации изображения
def generate_image(prompt, negative_prompt, model):
    if model == "CLIP":
        from clip import clip
        image = clip.generate(prompt, negative_prompt)
    elif model == "DALL-E 2":
        from d2l import d2l
        image = d2l.generate_image(prompt, negative_prompt)
    elif model == "VQGAN + CLIP":
        from vqgan_clip import vqgan_clip
        image = vqgan_clip.generate(prompt, negative_prompt)
    elif model == "VQGAN":
        from vqgan import vqgan
        image = vqgan.generate(prompt, negative_prompt)
    elif model == "Dreams":
        from dreams import dreams
        image = dreams.generate(prompt, negative_prompt)
    elif model == "VQGAN-CLIP-ADA":
        from vqgan_clip_ada import vqgan_clip_ada
        image = vqgan_clip_ada.generate(prompt, negative_prompt)
    elif model == "DALL-E Mini":
        from dall_e_mini import dall_e_mini
        image = dall_e_mini.generate(prompt, negative_prompt)
    elif model == "VQGAN-CLIP-Ada-Mini":
        from vqgan_clip_ada_mini import vqgan_clip_ada_mini
        image = vqgan_clip_ada_mini.generate(prompt, negative_prompt)
    elif model == "DALL-E Mini-Fusion":
        from dall_e_mini_fusion import dall_e_mini_fusion
        image = dall_e_mini_fusion.generate(prompt, negative_prompt)
    elif model == "VQGAN-CLIP-Ada-Mini-Fusion":
        from vqgan_clip_ada_mini_fusion import vqgan_clip_ada_mini_fusion
        image = vqgan_clip_ada_mini_fusion.generate(prompt, negative_prompt)
    elif model == "DALL-E 2 Mini":
        from dall_e_2_mini import dall_e_2_mini
        image = dall_e_2_mini.generate(prompt, negative_prompt)
    elif model == "VQGAN-CLIP-Ada-Mini-2":
        from vqgan_clip_ada_mini_2 import vqgan_clip_ada_mini_2
        image = vqgan_clip_ada_mini_2.generate(prompt, negative_prompt)

    image = Image.fromarray(image)
    return image

# Функция улучшения качества изображения
def improve_image(image):
    from improve_image import improve_image
    image = improve_image(image)
    return image

# Интерфейс gradio
app = gr.Interface(
    generate_image,
    title="Gradio Image Generator",
    description="Генератор изображений с помощью ИИ",
    inputs=[
        gr.Textbox(label="Текстовый запрос"),
        gr.Textbox(label="Отрицательный текстовый запрос"),
        gr.Dropdown(choices=models, label="Модель"),
    ],
    outputs=[
        gr.Image(label="Сгенерированное изображение")
    ],
    buttons=[
        gr.Button("Улучшить качество", fn=improve_image)
    ]
)

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