Update app.py

app.py

@@ -1,6 +1,6 @@
 import numpy as np
 from PIL import Image
-from huggingface_hub import snapshot_download
+from huggingface_hub import snapshot_download, login
 from leffa.transform import LeffaTransform
 from leffa.model import LeffaModel
 from leffa.inference import LeffaInference
@@ -13,11 +13,19 @@ from diffusers import DiffusionPipeline
 from transformers import pipeline
 import gradio as gr
 import os
-from huggingface_hub import login
 import random
+import gc
+
+# 메모리 최적화 설정
+torch.backends.cudnn.benchmark = True
+torch.backends.cuda.matmul.allow_tf32 = True
+os.environ['PYTORCH_CUDA_ALLOC_CONF'] = 'max_split_size_mb:512'
 
 # 상수 정의
 MAX_SEED = 2**32 - 1
+BASE_MODEL = "black-forest-labs/FLUX.1-dev"
+MODEL_LORA_REPO = "Motas/Flux_Fashion_Photography_Style"
+CLOTHES_LORA_REPO = "prithivMLmods/Canopus-Clothing-Flux-LoRA"
 
 # Hugging Face 토큰 설정 및 로그인
 HF_TOKEN = os.getenv("HF_TOKEN")
@@ -25,20 +33,26 @@ if HF_TOKEN is None:
     raise ValueError("Please set the HF_TOKEN environment variable")
 login(token=HF_TOKEN)
 
-# 모델 설정 (한 번만 선언)
-base_model = "black-forest-labs/FLUX.1-dev"
-model_lora_repo = "Motas/Flux_Fashion_Photography_Style"
-clothes_lora_repo = "prithivMLmods/Canopus-Clothing-Flux-LoRA"
+# 메모리 정리 함수
+def clear_memory():
+    torch.cuda.empty_cache()
+    gc.collect()
+
+# 초기 메모리 정리
+clear_memory()
 
-# FLUX 모델 초기화 (한 번만 초기화)
+# CUDA 사용 가능 여부 확인
+device = "cuda" if torch.cuda.is_available() else "cpu"
+
+# FLUX 모델 초기화
 fashion_pipe = DiffusionPipeline.from_pretrained(
-    base_model,
-    torch_dtype=torch.bfloat16,
+    BASE_MODEL,
+    torch_dtype=torch.float16,
     use_auth_token=HF_TOKEN
 )
-fashion_pipe.to("cuda")
+fashion_pipe.enable_model_cpu_offload()
 
-# 번역기 초기화 (한 번만)
+# 번역기 초기화
 translator = pipeline("translation", model="Helsinki-NLP/opus-mt-ko-en")
 
 # Leffa 체크포인트 다운로드
@@ -55,18 +69,22 @@ densepose_predictor = DensePosePredictor(
     weights_path="./ckpts/densepose/model_final_162be9.pkl",
 )
 
+# Leffa 모델 초기화
 vt_model = LeffaModel(
     pretrained_model_name_or_path="./ckpts/stable-diffusion-inpainting",
     pretrained_model="./ckpts/virtual_tryon.pth",
+    device_map="auto"
 )
 vt_inference = LeffaInference(model=vt_model)
 
 pt_model = LeffaModel(
     pretrained_model_name_or_path="./ckpts/stable-diffusion-xl-1.0-inpainting-0.1",
     pretrained_model="./ckpts/pose_transfer.pth",
+    device_map="auto"
 )
 pt_inference = LeffaInference(model=pt_model)
 
+
 def contains_korean(text):
     return any(ord('가') <= ord(char) <= ord('힣') for char in text)