Fix

app.py

@@ -19,12 +19,6 @@ load_dotenv()
 TOKEN = os.getenv("HF_TOKEN")
 login(token=TOKEN)
 
-torch.backends.cuda.matmul.allow_tf32 = True
-torch.backends.cudnn.benchmark = True
-torch.set_grad_enabled(False)
-os.environ["PYTORCH_CUDA_ALLOC_CONF"] = "max_split_size_mb:128"
-os.environ["CUDA_MODULE_LOADING"] = "LAZY"
-
 device = "cuda" if torch.cuda.is_available() else "cpu"
 weight_dtype = torch.float16 if device == "cuda" else torch.float32
 
@@ -104,7 +98,7 @@ def load_models():
         print(f"❌ Ошибка загрузки моделей: {e}")
         raise
 
-def generate_vton(person_image, cloth_image, outfit_prompt="", clothing_prompt=""):
+def generate_vton(person_image, cloth_image, outfit_prompt="", clothing_prompt="", label=7):
     """Генерация виртуальной примерки с очисткой памяти"""
     try:
         torch.cuda.empty_cache()
@@ -116,7 +110,7 @@ def generate_vton(person_image, cloth_image, outfit_prompt="", clothing_prompt="
             person_image.save(person_path)
             cloth_image.save(cloth_path)
             
-            mask_image = generate_clothing_mask(person_path)
+            mask_image = generate_clothing_mask(person_path, label=label)
             pose_image = generate_openpose(person_path)
             
             final_outfit_prompt = outfit_prompt or generate_caption(person_path, device)
@@ -158,37 +152,51 @@ pipeline = models["pipeline"]
 with gr.Blocks(theme=gr.themes.Soft(), css=".gradio-container") as demo:
     gr.Markdown("# 🧥 Virtual Try-On")
     gr.Markdown("Загрузите фото человека и одежды для виртуальной примерки")
-    
+
+    clothing_classes = [
+        "фон", "шляпа", "волосы", "сумка", "рюкзак", "верхняя одежда",
+        "рубашка", "футболка", "брюки", "короткие штаны", "платье", 
+        "юбка", "носки", "обувь", "водолазка", "аксессуары", "руки", "ноги"
+    ]
+
     with gr.Row():
         with gr.Column():
             person_input = gr.Image(label="Фото человека", type="pil", sources=["upload"])
             cloth_input = gr.Image(label="Фото одежды", type="pil", sources=["upload"])
+            clothing_label = gr.Dropdown(
+                choices=[(f"{i}: {desc}", i) for i, desc in enumerate(clothing_classes)],
+                label="Класс одежды для маски",
+                value=7
+            )
+
             outfit_prompt = gr.Textbox(label="Описание образа (опционально)", placeholder="Например: man in casual outfit")
             clothing_prompt = gr.Textbox(label="Описание одежды (опционально)", placeholder="Например: red t-shirt with print")
+
             generate_btn = gr.Button("Сгенерировать примерку", variant="primary")
-            
+
             gr.Examples(
                 examples=[
-                    ["./test/person2.png", "./test/00008_00.jpg", "man in skirt", "black longsleeve"]
+                    ["./test/person2.png", "./test/00008_00.jpg", "man in skirt", "black longsleeve", 7]
                 ],
-                inputs=[person_input, cloth_input, outfit_prompt, clothing_prompt],
+                inputs=[person_input, cloth_input, outfit_prompt, clothing_prompt, clothing_label],
                 label="Примеры для быстрого тестирования"
             )
-            
+
         with gr.Column():
             output_image = gr.Image(label="Результат примерки", interactive=False)
-    
+
     generate_btn.click(
         fn=generate_vton,
-        inputs=[person_input, cloth_input, outfit_prompt, clothing_prompt],
+        inputs=[person_input, cloth_input, outfit_prompt, clothing_prompt, clothing_label],
         outputs=output_image
     )
-    
+
     gr.Markdown("### Инструкция:")
     gr.Markdown("1. Загрузите четкое фото человека в полный рост\n"
                 "2. Загрузите фото одежды на белом фоне\n"
-                "3. При необходимости уточните описание образа или одежды\n"
-                "4. Нажмите кнопку 'Сгенерировать примерку'")
+                "3. Выберите тип одежды из выпадающего списка\n"
+                "4. При необходимости уточните описание образа или одежды\n"
+                "5. Нажмите кнопку 'Сгенерировать примерку'")
 
 if __name__ == "__main__":
     demo.queue(max_size=1).launch(

lib/mask.py

@@ -9,9 +9,8 @@ import os
 def generate_clothing_mask(
     image_path: str,
     label: int,
-    output_path: str = "./output_mask.png",
+    output_path: str = "./test/output_mask.png",
     model_name: str = "mattmdjaga/segformer_b2_clothes",
-    show_result: bool = False
 ) -> Image.Image:
     """
     Генерирует бинарную маску для указанного класса одежды и сохраняет её