Update app.py

app.py

@@ -172,29 +172,33 @@ def get_densepose_predictor():
         )
     return densepose_predictor
 
-@safe_model_call
+@spaces.GPU()
 def get_vt_model():
-    global vt_model, vt_inference
-    if vt_model is None:
-        vt_model = LeffaModel(
+    try:
+        model = LeffaModel(
             pretrained_model_name_or_path="./ckpts/stable-diffusion-inpainting",
             pretrained_model="./ckpts/virtual_tryon.pth"
         )
-        vt_model = vt_model.half().to(device)
-        vt_inference = LeffaInference(model=vt_model)
-    return vt_model, vt_inference
+        model = model.half().to("cuda")
+        inference = LeffaInference(model=model)
+        return model, inference
+    except Exception as e:
+        print(f"Error in get_vt_model: {str(e)}")
+        raise
 
-@safe_model_call
+@spaces.GPU()
 def get_pt_model():
-    global pt_model, pt_inference
-    if pt_model is None:
-        pt_model = LeffaModel(
+    try:
+        model = LeffaModel(
             pretrained_model_name_or_path="./ckpts/stable-diffusion-xl-1.0-inpainting-0.1",
             pretrained_model="./ckpts/pose_transfer.pth"
         )
-        pt_model = pt_model.half().to(device)
-        pt_inference = LeffaInference(model=pt_model)
-    return pt_model, pt_inference
+        model = model.half().to("cuda")
+        inference = LeffaInference(model=model)
+        return model, inference
+    except Exception as e:
+        print(f"Error in get_pt_model: {str(e)}")
+        raise
 
 def load_lora(pipe, lora_path):
     try:
@@ -298,7 +302,7 @@ def generate_fashion(prompt, mode, cfg_scale, steps, randomize_seed, seed, width
         print(f"Error in generate_fashion: {str(e)}")
         raise gr.Error(f"Generation failed: {str(e)}")
 
-@safe_model_call
+@spaces.GPU()
 def leffa_predict(src_image_path, ref_image_path, control_type):
     try:
         # 모델 초기화
@@ -307,10 +311,7 @@ def leffa_predict(src_image_path, ref_image_path, control_type):
         else:
             model, inference = get_pt_model()
 
-        mask_pred = get_mask_predictor()
-        dense_pred = get_densepose_predictor()
-
-        # 이미지 로드 및 전처리
+        # 이미지 처리
         src_image = Image.open(src_image_path)
         ref_image = Image.open(ref_image_path)
         src_image = resize_and_center(src_image, 768, 1024)
@@ -319,21 +320,23 @@ def leffa_predict(src_image_path, ref_image_path, control_type):
         src_image_array = np.array(src_image)
         ref_image_array = np.array(ref_image)
 
-        # Mask 생성
-        if control_type == "virtual_tryon":
-            src_image = src_image.convert("RGB")
-            mask = mask_pred(src_image, "upper")["mask"]
-        else:
-            mask = Image.fromarray(np.ones_like(src_image_array) * 255)
-
-        # DensePose 예측
-        src_image_iuv_array = dense_pred.predict_iuv(src_image_array)
-        src_image_seg_array = dense_pred.predict_seg(src_image_array)
-        
-        if control_type == "virtual_tryon":
-            densepose = Image.fromarray(src_image_seg_array)
-        else:
-            densepose = Image.fromarray(src_image_iuv_array)
+        # Mask 및 DensePose 처리
+        with torch.inference_mode():
+            if control_type == "virtual_tryon":
+                src_image = src_image.convert("RGB")
+                mask_pred = get_mask_predictor()
+                mask = mask_pred(src_image, "upper")["mask"]
+            else:
+                mask = Image.fromarray(np.ones_like(src_image_array) * 255)
+
+            dense_pred = get_densepose_predictor()
+            src_image_iuv_array = dense_pred.predict_iuv(src_image_array)
+            src_image_seg_array = dense_pred.predict_seg(src_image_array)
+            
+            if control_type == "virtual_tryon":
+                densepose = Image.fromarray(src_image_seg_array)
+            else:
+                densepose = Image.fromarray(src_image_iuv_array)
 
         # Leffa 변환 및 추론
         transform = LeffaTransform()
@@ -345,20 +348,36 @@ def leffa_predict(src_image_path, ref_image_path, control_type):
         }
         data = transform(data)
         
-        output = inference(data)
+        with torch.inference_mode():
+            output = inference(data)
+
+        # 메모리 정리
+        del model
+        del inference
+        torch.cuda.empty_cache()
+        gc.collect()
+
         return np.array(output["generated_image"][0])
         
     except Exception as e:
         print(f"Error in leffa_predict: {str(e)}")
         raise
 
-@safe_model_call
+@spaces.GPU()
 def leffa_predict_vt(src_image_path, ref_image_path):
-    return leffa_predict(src_image_path, ref_image_path, "virtual_tryon")
+    try:
+        return leffa_predict(src_image_path, ref_image_path, "virtual_tryon")
+    except Exception as e:
+        print(f"Error in leffa_predict_vt: {str(e)}")
+        raise
 
-@safe_model_call
+@spaces.GPU()
 def leffa_predict_pt(src_image_path, ref_image_path):
-    return leffa_predict(src_image_path, ref_image_path, "pose_transfer")
+    try:
+        return leffa_predict(src_image_path, ref_image_path, "pose_transfer")
+    except Exception as e:
+        print(f"Error in leffa_predict_pt: {str(e)}")
+        raise
 
 # 초기 설정 실행
 setup()