Update app.py

app.py

@@ -32,6 +32,20 @@ def clear_memory():
         torch.cuda.synchronize()
     gc.collect()
 
+# GPU 메모리 관리 설정
+if torch.cuda.is_available():
+    torch.cuda.empty_cache()
+    torch.backends.cudnn.benchmark = True
+    torch.backends.cuda.matmul.allow_tf32 = True
+    
+    # 메모리 분할 설정
+    os.environ["PYTORCH_CUDA_ALLOC_CONF"] = (
+        "max_split_size_mb:128,"
+        "garbage_collection_threshold:0.8,"
+        "memory_fraction:0.9"
+    )
+
+    
 # 자동 혼합 정밀도(Automatic Mixed Precision) 설정
 if torch.cuda.is_available():
     scaler = torch.amp.GradScaler('cuda')
@@ -89,6 +103,7 @@ gd_model = gd_model.to(device=device)
 assert isinstance(gd_model, GroundingDinoForObjectDetection)
 
 
+# FLUX 파이프라인 초기화
 # FLUX 파이프라인 초기화
 pipe = FluxPipeline.from_pretrained(
     "black-forest-labs/FLUX.1-dev",
@@ -110,9 +125,9 @@ pipe.fuse_lora(lora_scale=0.125)
 # GPU 메모리 최적화
 if torch.cuda.is_available():
     pipe.to("cuda")
-    pipe.enable_vae_slicing()  # VAE 슬라이싱 활성화
-    
-
+    # 메모리 최적화를 위한 추가 설정
+    torch.backends.cudnn.benchmark = True
+    torch.backends.cuda.matmul.allow_tf32 = True
 
 os.environ["CUDA_VISIBLE_DEVICES"] = "0"  # 단일 GPU 사용
 os.environ["PYTORCH_CUDA_ALLOC_CONF"] = "max_split_size_mb:512"  # CUDA 메모리 할당 설정
@@ -203,40 +218,26 @@ def generate_background(prompt: str, aspect_ratio: str) -> Image.Image:
             height = int(height * ratio)
             width, height = adjust_size_to_multiple_of_8(width, height)
 
+        clear_memory()  # 생성 전 메모리 정리
+        
         with timer("Background generation"):
-            with torch.inference_mode():  # inference_mode 사용
-                with torch.cuda.amp.autocast():
-                    image = pipe(
-                        prompt=prompt,
-                        width=width,
-                        height=height,
-                        num_inference_steps=8,
-                        guidance_scale=4.0,
-                        max_length=77,
-                    ).images[0]
-                    
-            clear_memory()  # 즉시 메모리 정리
-            return image
+            with torch.inference_mode(), torch.cuda.amp.autocast():
+                image = pipe(
+                    prompt=prompt,
+                    width=width,
+                    height=height,
+                    num_inference_steps=8,
+                    guidance_scale=4.0,
+                    max_length=77,
+                ).images[0]
+        
+        clear_memory()  # 생성 후 메모리 정리
+        return image
     except Exception as e:
         print(f"Background generation error: {str(e)}")
+        clear_memory()  # 오류 발생 시에도 메모리 정리
         return Image.new('RGB', (512, 512), 'white')
 
-# FLUX 파이프라인 초기화 부분 수정
-pipe = FluxPipeline.from_pretrained(
-    "black-forest-labs/FLUX.1-dev",
-    torch_dtype=torch.float32,  # bfloat16 대신 float32 사용
-    use_auth_token=HF_TOKEN
-)
-pipe.load_lora_weights(
-    hf_hub_download(
-        "ByteDance/Hyper-SD",
-        "Hyper-FLUX.1-dev-8steps-lora.safetensors",
-        use_auth_token=HF_TOKEN
-    )
-)
-pipe.fuse_lora(lora_scale=0.125)
-pipe.to(device=device)
-
 
 def create_position_grid():
     return """