Update app.py

app.py

@@ -36,21 +36,17 @@ def safe_model_call(func):
             raise
     return wrapper
 
-# 메모리 관리 함수
+# 메모리 관리 함수 수정
 def clear_memory():
-    if torch.cuda.is_available():
-        torch.cuda.empty_cache()
-        torch.cuda.synchronize()
     gc.collect()
+    if torch.cuda.is_available() and torch.cuda.current_device() >= 0:
+        torch.cuda.empty_cache()
 
+# 환경 설정 함수 수정
 def setup_environment():
     # 메모리 관리 설정
-    torch.cuda.empty_cache()
     gc.collect()
     os.environ['PYTORCH_CUDA_ALLOC_CONF'] = 'max_split_size_mb:128'
-    torch.backends.cudnn.benchmark = True
-    torch.backends.cuda.matmul.allow_tf32 = True
-    torch.backends.cuda.max_split_size_mb = 128
 
     # Hugging Face 토큰 설정
     global HF_TOKEN
@@ -59,9 +55,11 @@ def setup_environment():
         raise ValueError("Please set the HF_TOKEN environment variable")
     login(token=HF_TOKEN)
 
-    # CUDA 설정
+    # device 설정 제거 (spaces.GPU() 데코레이터가 처리)
     global device
-    device = "cuda" if torch.cuda.is_available() else "cpu"
+    device = "cpu"  # 기본값으로 CPU 설정
+
+
 
 # 전역 변수 초기화
 fashion_pipe = None
@@ -78,31 +76,32 @@ HF_TOKEN = None
 # 환경 설정 실행
 setup_environment()
 
-
-# 모델 관리 함수들
+@spaces.GPU()
 def initialize_fashion_pipe():
     global fashion_pipe
     if fashion_pipe is None:
-        clear_memory()
         fashion_pipe = DiffusionPipeline.from_pretrained(
             BASE_MODEL,
             torch_dtype=torch.float16,
             use_auth_token=HF_TOKEN
-        )
+        ).to("cuda")
         try:
             fashion_pipe.enable_xformers_memory_efficient_attention()
         except Exception as e:
             print(f"Warning: Could not enable memory efficient attention: {e}")
-        fashion_pipe.enable_sequential_cpu_offload()
     return fashion_pipe
 
-@safe_model_call
+def setup():
+    # Leffa 체크포인트 다운로드만 수행
+    snapshot_download(repo_id="franciszzj/Leffa", local_dir="./ckpts")
+    
+@spaces.GPU()
 def get_translator():
     global translator
     if translator is None:
         translator = pipeline("translation", 
                             model="Helsinki-NLP/opus-mt-ko-en",
-                            device=device if device == "cuda" else -1)
+                            device="cuda")
     return translator
 
 @safe_model_call
@@ -161,13 +160,16 @@ def load_lora(pipe, lora_path):
         print(f"Warning: Failed to load LoRA weights from {lora_path}: {e}")
         return pipe
 
-# 초기 설정 함수
-def setup():
-    # Leffa 체크포인트 다운로드
-    snapshot_download(repo_id="franciszzj/Leffa", local_dir="./ckpts")
-    # 기본 모델 초기화
-    initialize_fashion_pipe()
-
+@spaces.GPU()
+def get_mask_predictor():
+    global mask_predictor
+    if mask_predictor is None:
+        mask_predictor = AutoMasker(
+            densepose_path="./ckpts/densepose",
+            schp_path="./ckpts/schp",
+        )
+    return mask_predictor
+    
 # 유틸리티 함수
 def contains_korean(text):
     return any(ord('가') <= ord(char) <= ord('힣') for char in text)
@@ -175,7 +177,6 @@ def contains_korean(text):
 
 # 메인 기능 함수들
 @spaces.GPU()
-@safe_model_call
 def generate_fashion(prompt, mode, cfg_scale, steps, randomize_seed, seed, width, height, lora_scale, progress=gr.Progress(track_tqdm=True)):
     try:
         # 한글 처리