Update app.py

app.py

@@ -1,20 +1,22 @@
 import gradio as gr
 import numpy as np
-import random
 import torch
 from diffusers import DiffusionPipeline
+import re
 
+# Устройство и параметры загрузки модели
 device = "cuda" if torch.cuda.is_available() else "cpu"
-model_repo_id = "CompVis/stable-diffusion-v1-4"  # Replace to the model you would like to use
+torch_dtype = torch.float16 if torch.cuda.is_available() else torch.float32
 
-if torch.cuda.is_available():
-    torch_dtype = torch.float16
-else:
-    torch_dtype = torch.float32
+# Регулярное выражение для проверки корректности модели
+VALID_REPO_ID_REGEX = re.compile(r"^[a-zA-Z0-9._\-]+/[a-zA-Z0-9._\-]+$")
 
-# Изначально загружаем модель по умолчанию (как в исходном коде)
-pipe = DiffusionPipeline.from_pretrained(model_repo_id, torch_dtype=torch_dtype)
-pipe = pipe.to(device)
+def is_valid_repo_id(repo_id):
+    return bool(VALID_REPO_ID_REGEX.match(repo_id)) and not repo_id.endswith(('-', '.'))
+
+# Изначально загружаем модель по умолчанию
+model_repo_id = "CompVis/stable-diffusion-v1-4"
+pipe = DiffusionPipeline.from_pretrained(model_repo_id, torch_dtype=torch_dtype).to(device)
 
 MAX_SEED = np.iinfo(np.int32).max
 MAX_IMAGE_SIZE = 1024
@@ -31,28 +33,34 @@ def infer(
     progress=gr.Progress(track_tqdm=True),
 ):
     global model_repo_id, pipe
-    
-    # Если пользователь ввёл новую модель, пробуем загрузить
+
+    # Проверяем и загружаем новую модель, если она изменена
     if model != model_repo_id:
+        if not is_valid_repo_id(model):
+            raise gr.Error(f"Некорректный идентификатор модели: '{model}'. Проверьте название.")
         try:
-            new_pipe = DiffusionPipeline.from_pretrained(model, torch_dtype=torch_dtype)
-            new_pipe = new_pipe.to(device)
+            new_pipe = DiffusionPipeline.from_pretrained(model, torch_dtype=torch_dtype).to(device)
             pipe = new_pipe
             model_repo_id = model
         except Exception as e:
             raise gr.Error(f"Не удалось загрузить модель '{model}'.\nОшибка: {e}")
 
+    # Генератор случайных чисел для детерминированности
     generator = torch.Generator(device=device).manual_seed(seed)
 
-    image = pipe(
-        prompt=prompt,
-        negative_prompt=negative_prompt,
-        guidance_scale=guidance_scale,
-        num_inference_steps=num_inference_steps,
-        width=width,
-        height=height,
-        generator=generator,
-    ).images[0]
+    # Генерация изображения
+    try:
+        image = pipe(
+            prompt=prompt,
+            negative_prompt=negative_prompt,
+            guidance_scale=guidance_scale,
+            num_inference_steps=num_inference_steps,
+            width=width,
+            height=height,
+            generator=generator,
+        ).images[0]
+    except Exception as e:
+        raise gr.Error(f"Ошибка при генерации изображения: {e}")
 
     return image, seed
 
@@ -69,14 +77,13 @@ css = """
 }
 """
 
-# Убрали выпадающее меню, заменили на текстовое поле
 with gr.Blocks(css=css) as demo:
     with gr.Column(elem_id="col-container"):
-        gr.Markdown(" # Text-to-Image App")
+        gr.Markdown("# Text-to-Image App")
 
         model = gr.Textbox(
             label="Model",
-            value="stabilityai/sdxl-turbo",  # Значение по умолчанию
+            value="CompVis/stable-diffusion-v1-4",  # Значение по умолчанию
             interactive=True
         )
 
@@ -87,14 +94,14 @@ with gr.Blocks(css=css) as demo:
             placeholder="Enter your prompt",
             container=False,
         )
-        
+
         negative_prompt = gr.Text(
             label="Negative prompt",
             max_lines=1,
             placeholder="Enter a negative prompt",
             visible=True,
         )
-        
+
         seed = gr.Slider(
             label="Seed",
             minimum=0,
@@ -102,7 +109,7 @@ with gr.Blocks(css=css) as demo:
             step=1,
             value=42,
         )
-        
+
         guidance_scale = gr.Slider(
             label="Guidance scale",
             minimum=0.0,
@@ -118,8 +125,8 @@ with gr.Blocks(css=css) as demo:
             step=1,
             value=20,
         )
-        
-        run_button = gr.Button("Run", scale=0, variant="primary")
+
+        run_button = gr.Button("Run", variant="primary")
         result = gr.Image(label="Result", show_label=False)
 
         with gr.Accordion("Advanced Settings", open=False):
@@ -129,33 +136,32 @@ with gr.Blocks(css=css) as demo:
                     minimum=256,
                     maximum=MAX_IMAGE_SIZE,
                     step=32,
-                    value=1024,
+                    value=512,
                 )
                 height = gr.Slider(
                     label="Height",
                     minimum=256,
                     maximum=MAX_IMAGE_SIZE,
                     step=32,
-                    value=1024,
+                    value=512,
                 )
 
         gr.Examples(examples=examples, inputs=[prompt])
 
-    gr.on(
-        triggers=[run_button.click, prompt.submit],
-        fn=infer,
-        inputs=[
-            model,
-            prompt,
-            negative_prompt,
-            seed,
-            width,
-            height,
-            guidance_scale,
-            num_inference_steps,
-        ],
-        outputs=[result, seed],
-    )
+        run_button.click(
+            infer,
+            inputs=[
+                model,
+                prompt,
+                negative_prompt,
+                seed,
+                width,
+                height,
+                guidance_scale,
+                num_inference_steps,
+            ],
+            outputs=[result, seed],
+        )
 
 if __name__ == "__main__":
     demo.launch()