Update app.py

app.py

@@ -1,51 +1,84 @@
 import gradio as gr
-import requests
-import os
+import torch
+from PIL import Image
 
-# Получаем URL и API Key из переменных окружения
-base_url = os.getenv("BASE_URL")
-api_key = os.getenv("API_KEY")
+# Список моделей
+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, sampling_steps, cfg_scale, seed):
-    headers = {
-        "Authorization": f"Bearer {api_key}"
-    }
+# Функция генерации изображения
+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)
 
-    payload = {
-        "prompt": prompt,
-        "negative_prompt": negative_prompt,
-        "sampling_steps": sampling_steps,
-        "cfg_scale": cfg_scale,
-        "seed": seed
-    }
+    image = Image.fromarray(image)
+    return image
 
-    response = requests.post(base_url, json=payload, headers=headers)
+# Функция улучшения качества изображения
+def improve_image(image):
+    from improve_image import improve_image
+    image = improve_image(image)
+    return image
 
-    # Предположим, что модель возвращает изображение в виде байтов
-    generated_image_bytes = response.content
-
-    return generated_image_bytes
-
-# Функция для обработки действия при нажатии кнопки
-def download_image(prompt, negative_prompt, sampling_steps, cfg_scale, seed):
-    bytes = generate_image(prompt, negative_prompt, sampling_steps, cfg_scale, seed)
-    gr.download(bytes, "generated_image.png")
-
-# Создаем интерфейс Gradio
-iface = gr.Interface(
-    fn=download_image,
-    inputs=[
-        gr.Textbox(label="Prompt", placeholder="Enter prompt here"),
-        gr.Textbox(label="Negative Prompt", placeholder="Enter negative prompt here"),
-        gr.Slider(minimum=1, maximum=30, label="Sampling Steps", initial=15),
-        gr.Slider(minimum=1, maximum=20, label="CFG Scale", initial=10),
-        gr.Textbox(label="Seed", placeholder="Enter seed here", default="-1")
-    ],
-    outputs=gr.Image(type="pil", label="Generated Image"),
-    live=True,
-    title="Huggingface Image Generator"
+# Интерфейс gradio
+app = gr.Interface(
+    generate_image,
+    title="Gradio Image Generator",
+    description="Генератор изображений с помощью ИИ",
+    inputs={
+        "prompt": gr.Input(type="text", label="Текстовый запрос"),
+        "negative_prompt": gr.Input(type="text", label="Отрицательный текстовый запрос"),
+        "model": gr.Select(models, label="Модель"),
+    },
+    outputs={"image": gr.Output(type="image", label="Сгенерированное изображение")},
+    buttons={"Улучшить качество": improve_image},
 )
 
-# Запускаем интерфейс
-iface.launch()
+if __name__ == "__main__":
+    app.launch()