optimization

app.py

@@ -20,8 +20,22 @@ print(f"GPU count: {torch.cuda.device_count()}")
 if torch.cuda.is_available():
     print(f"Current device: {torch.cuda.current_device()}")
     print(f"Device name: {torch.cuda.get_device_name(0)}")
+    # Настройка PyTorch для A100
+    torch.backends.cuda.enable_flash_sdp(True)  # Включение Flash Attention
+    torch.backends.cuda.enable_mem_efficient_sdp(True)  # Экономия памяти
+    torch.set_float32_matmul_precision('high')  # Оптимизация матричных операций
 print("="*50)
 
+def clear_cuda():
+    if torch.cuda.is_available():
+        print(f"Используется VRAM: {torch.cuda.memory_allocated() / 1024 ** 3:.2f} GB")
+        print(f"Доступно VRAM: {torch.cuda.memory_reserved() / 1024 ** 3:.2f} GB")
+        print(f"Очистка кеша CUDA...")
+        torch.cuda.empty_cache()
+        print(f"Очистка кеша CUDA завершена.")
+        print(f"Используется VRAM: {torch.cuda.memory_allocated() / 1024 ** 3:.2f} GB")
+        print(f"Доступно VRAM: {torch.cuda.memory_reserved() / 1024 ** 3:.2f} GB")
+
 device = "cuda" if torch.cuda.is_available() else "cpu"
 
 print(f"{datetime.datetime.now()} Загрузка модели FLUX.1-dev")
@@ -43,6 +57,8 @@ print(f"{datetime.datetime.now()} Загрузка LoRA успешно заве
 
 pipe.fuse_lora(lora_scale=1.0)
 pipe.to(device)
+
+pipe.enable_xformers_memory_efficient_attention()  # Ускорение внимания
 pipe.enable_model_cpu_offload()  # Выгрузка неиспользуемых компонентов
 
 # print(f"{datetime.datetime.now()} Загрузка модели stabilityai/stable-diffusion-x4-upscaler")
@@ -58,6 +74,8 @@ print(f"{datetime.datetime.now()} Загрузка модели briaai/RMBG-1.4"
 bg_remover = pipeline("image-segmentation", "briaai/RMBG-1.4", trust_remote_code=True )
 print(f"{datetime.datetime.now()} Загрузка модели briaai/RMBG-1.4 успешно завершена")
 
+clear_cuda()
+
 @spaces.GPU()
 def generate_image(object_name, remove_bg=True):
     try:
@@ -73,17 +91,18 @@ def generate_image(object_name, remove_bg=True):
         steps = os.getenv('STEPS') if os.getenv('STEPS') is not None else 10
         print(f"Шаги: {steps}")
 
+        clear_cuda()
+
         image = pipe(
             prompt,
-            height=1024,
-            width=1024,
-            guidance_scale=3.5,
+            height=768,
+            width=768,
+            guidance_scale=4.0,
             num_inference_steps=int(steps),
-            generator=torch.Generator(device).manual_seed(42)
+            generator=torch.Generator(device).manual_seed(42),
+            num_images_per_prompt=1
         ).images[0]
 
-        torch.cuda.empty_cache()
-
         # if upscale :
         #     torch.cuda.empty_cache()
         #     upscaled_image = upscaler_pipeline(
@@ -95,11 +114,12 @@ def generate_image(object_name, remove_bg=True):
         #     return upscaled_image
 
         if remove_bg :
+            clear_cuda()
             remove_bg_image = bg_remover(image)
-            torch.cuda.empty_cache()
+            clear_cuda()
             return remove_bg_image
 
-        torch.cuda.empty_cache()
+        clear_cuda()
         return image
     
     except Exception as e: