Update app.py

app.py

@@ -1,13 +1,24 @@
 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("cuda")
+    models[name] = pipeline
 
 # Функция для генерации изображения
 def generate_image(prompt, model_name, sampling_steps, cfg, sampling_method, seed, width, height):
-    # Установка модели в зависимости от выбранного названия
-    if model_name == "Vector Art":
-        model_path = "stablediffusionapi/vector-art"
-    # Загрузка модели (пример с использованием StableDiffusionPipeline)
-    model = StableDiffusionPipeline.from_pretrained(model_path)
+    # Выбор модели из словаря
+    model = models[model_name]
     # Установка параметров генерации
     generator = torch.manual_seed(seed) if seed != -1 else None
     # Генерация изображения
@@ -21,7 +32,7 @@ with gr.Blocks() as app:
     with gr.Tab("Основные настройки"):
         with gr.Row():
             prompt_input = gr.Textbox(label="Введите текстовый запрос")
-            model_select = gr.Radio(choices=["Vector Art"], label="Выберите модель", value="Vector Art")
+            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)