Update app.py

app.py

@@ -2,9 +2,10 @@ import gradio as gr
 import numpy as np
 import torch
 from diffusers import DiffusionPipeline
+from peft import PeftModel
 import re
 
-# Устройство и параметры загрузки модели
+# Устройство и тип данных
 device = "cuda" if torch.cuda.is_available() else "cpu"
 torch_dtype = torch.float16 if torch.cuda.is_available() else torch.float32
 
@@ -14,15 +15,21 @@ VALID_REPO_ID_REGEX = re.compile(r"^[a-zA-Z0-9._\-]+/[a-zA-Z0-9._\-]+$")
 def is_valid_repo_id(repo_id):
     return bool(VALID_REPO_ID_REGEX.match(repo_id)) and not repo_id.endswith(('-', '.'))
 
+# Базовые константы
+MAX_SEED = np.iinfo(np.int32).max
+MAX_IMAGE_SIZE = 1024
+
 # Изначально загружаем модель по умолчанию
 model_repo_id = "CompVis/stable-diffusion-v1-4"
 pipe = DiffusionPipeline.from_pretrained(model_repo_id, torch_dtype=torch_dtype).to(device)
 
-# --- Загрузка LoRA (1) ---
-pipe.load_lora_weights("AnastasiaSh/sticker-cat-lora3")
-
-MAX_SEED = np.iinfo(np.int32).max
-MAX_IMAGE_SIZE = 1024
+# Попробуем подгрузить LoRA-модификации (unet + text_encoder)
+try:
+    pipe.unet = PeftModel.from_pretrained(pipe.unet, "AnastasiaSh/sticker-cat-lora3/unet")
+    pipe.text_encoder = PeftModel.from_pretrained(pipe.text_encoder, "AnastasiaSh/sticker-cat-lora3/text_encoder")
+except Exception as e:
+    # Если не удалось, можно вывести предупреждение или поднять ошибку
+    print(f"Не удалось подгрузить LoRA по умолчанию: {e}")
 
 def infer(
     model,
@@ -37,25 +44,32 @@ def infer(
 ):
     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).to(device)
-            
-            # --- Загрузка LoRA (2) ---
-            new_pipe.load_lora_weights("AnastasiaSh/sticker-cat-lora3")
-            
+
+            # Повторно подгружаем LoRA для нового пайплайна
+            try:
+                new_pipe.unet = PeftModel.from_pretrained(new_pipe.unet, "AnastasiaSh/sticker-cat-lora3/unet")
+                new_pipe.text_encoder = PeftModel.from_pretrained(new_pipe.text_encoder, "AnastasiaSh/sticker-cat-lora3/text_encoder")
+            except Exception as e:
+                raise gr.Error(f"Не удалось подгрузить LoRA: {e}")
+
+            # Обновляем глобальные переменные
             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)
 
-    # Генерация изображения
+    # Пытаемся сгенерировать изображение
     try:
         image = pipe(
             prompt=prompt,
@@ -71,12 +85,14 @@ def infer(
 
     return image, seed
 
+# Примеры для удобного тестирования
 examples = [
     "Astronaut in a jungle, cold color palette, muted colors, detailed, 8k",
     "An astronaut riding a green horse",
     "A delicious ceviche cheesecake slice",
 ]
 
+# Дополнительный CSS для оформления
 css = """
 #col-container {
     margin: 0 auto;
@@ -84,16 +100,19 @@ css = """
 }
 """
 
+# Создаём Gradio-приложение
 with gr.Blocks(css=css) as demo:
     with gr.Column(elem_id="col-container"):
         gr.Markdown("# Text-to-Image App")
 
+        # Поле для ввода/смены модели
         model = gr.Textbox(
             label="Model",
             value="CompVis/stable-diffusion-v1-4",  # Значение по умолчанию
             interactive=True
         )
 
+        # Основные поля для Prompt и Negative Prompt
         prompt = gr.Text(
             label="Prompt",
             show_label=False,
@@ -101,7 +120,6 @@ with gr.Blocks(css=css) as demo:
             placeholder="Enter your prompt",
             container=False,
         )
-
         negative_prompt = gr.Text(
             label="Negative prompt",
             max_lines=1,
@@ -109,6 +127,7 @@ with gr.Blocks(css=css) as demo:
             visible=True,
         )
 
+        # Слайдер для выбора seed
         seed = gr.Slider(
             label="Seed",
             minimum=0,
@@ -117,6 +136,7 @@ with gr.Blocks(css=css) as demo:
             value=42,
         )
 
+        # Слайдеры для guidance_scale и num_inference_steps
         guidance_scale = gr.Slider(
             label="Guidance scale",
             minimum=0.0,
@@ -124,7 +144,6 @@ with gr.Blocks(css=css) as demo:
             step=0.1,
             value=7.0,
         )
-
         num_inference_steps = gr.Slider(
             label="Number of inference steps",
             minimum=1,
@@ -133,9 +152,13 @@ with gr.Blocks(css=css) as demo:
             value=20,
         )
 
+        # Кнопка запуска
         run_button = gr.Button("Run", variant="primary")
+
+        # Поле для отображения результата
         result = gr.Image(label="Result", show_label=False)
 
+        # Продвинутые настройки (Accordion)
         with gr.Accordion("Advanced Settings", open=False):
             with gr.Row():
                 width = gr.Slider(
@@ -153,8 +176,10 @@ with gr.Blocks(css=css) as demo:
                     value=512,
                 )
 
+        # Примеры
         gr.Examples(examples=examples, inputs=[prompt])
 
+        # Связка кнопки "Run" с функцией "infer"
         run_button.click(
             infer,
             inputs=[
@@ -170,5 +195,6 @@ with gr.Blocks(css=css) as demo:
             outputs=[result, seed],
         )
 
+# Запуск
 if __name__ == "__main__":
     demo.launch()