Update app.py

app.py

@@ -251,67 +251,105 @@ 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)}")
 
-@spaces.GPU()
-def leffa_predict(src_image_path, ref_image_path, control_type):
-    try:
-        # 모델 초기화
-        if control_type == "virtual_tryon":
-            model, inference = get_vt_model()
-        else:
-            model, inference = get_pt_model()
+class ModelManager:
+    def __init__(self):
+        self.mask_predictor = None
+        self.densepose_predictor = None
+        self.translator = None
+
+    @spaces.GPU()
+    def get_mask_predictor(self):
+        if self.mask_predictor is None:
+            self.mask_predictor = AutoMasker(
+                densepose_path="./ckpts/densepose",
+                schp_path="./ckpts/schp",
+            )
+        return self.mask_predictor
+
+    @spaces.GPU()
+    def get_densepose_predictor(self):
+        if self.densepose_predictor is None:
+            self.densepose_predictor = DensePosePredictor(
+                config_path="./ckpts/densepose/densepose_rcnn_R_50_FPN_s1x.yaml",
+                weights_path="./ckpts/densepose/model_final_162be9.pkl",
+            )
+        return self.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)
-        ref_image = resize_and_center(ref_image, 768, 1024)
+    @spaces.GPU()
+    def get_translator(self):
+        if self.translator is None:
+            self.translator = pipeline("translation", 
+                                    model="Helsinki-NLP/opus-mt-ko-en",
+                                    device="cuda")
+        return self.translator
 
-        src_image_array = np.array(src_image)
-        ref_image_array = np.array(ref_image)
+# 모델 매니저 인스턴스 생성
+model_manager = ModelManager()
 
-        # Mask 및 DensePose 처리
-        with torch.inference_mode():
+@spaces.GPU()
+def leffa_predict(src_image_path, ref_image_path, control_type):
+    try:
+        with torch_gc():
+            # 모델 초기화
             if control_type == "virtual_tryon":
-                src_image = src_image.convert("RGB")
-                mask_pred = get_mask_predictor()
-                mask = mask_pred(src_image, "upper")["mask"]
+                model, inference = get_vt_model()
             else:
-                mask = Image.fromarray(np.ones_like(src_image_array) * 255)
+                model, inference = get_pt_model()
+
+            # 이미지 처리
+            src_image = Image.open(src_image_path)
+            ref_image = Image.open(ref_image_path)
+            src_image = resize_and_center(src_image, 768, 1024)
+            ref_image = resize_and_center(ref_image, 768, 1024)
 
-            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)
+            src_image_array = np.array(src_image)
+            ref_image_array = np.array(ref_image)
+
+            # Mask 및 DensePose 처리
+            with torch.inference_mode():
+                if control_type == "virtual_tryon":
+                    src_image = src_image.convert("RGB")
+                    mask_pred = model_manager.get_mask_predictor()
+                    mask = mask_pred(src_image, "upper")["mask"]
+                else:
+                    mask = Image.fromarray(np.ones_like(src_image_array) * 255)
+
+                dense_pred = model_manager.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()
+            data = {
+                "src_image": [src_image],
+                "ref_image": [ref_image],
+                "mask": [mask],
+                "densepose": [densepose],
+            }
+            data = transform(data)
             
-            if control_type == "virtual_tryon":
-                densepose = Image.fromarray(src_image_seg_array)
-            else:
-                densepose = Image.fromarray(src_image_iuv_array)
-
-        # Leffa 변환 및 추론
-        transform = LeffaTransform()
-        data = {
-            "src_image": [src_image],
-            "ref_image": [ref_image],
-            "mask": [mask],
-            "densepose": [densepose],
-        }
-        data = transform(data)
-        
-        with torch.inference_mode():
-            output = inference(data)
+            with torch.inference_mode():
+                output = inference(data)
 
-        # 메모리 정리
-        del model
-        del inference
-        torch.cuda.empty_cache()
-        gc.collect()
+            # 메모리 정리
+            del model
+            del inference
+            torch.cuda.empty_cache()
+            gc.collect()
 
-        return np.array(output["generated_image"][0])
-        
+            return np.array(output["generated_image"][0])
+            
     except Exception as e:
         print(f"Error in leffa_predict: {str(e)}")
         raise
 
+
+        
 @spaces.GPU()
 def leffa_predict_vt(src_image_path, ref_image_path):
     try:
@@ -328,20 +366,20 @@ def leffa_predict_pt(src_image_path, ref_image_path):
         print(f"Error in leffa_predict_pt: {str(e)}")
         raise
 
-
 @spaces.GPU()
 def generate_image(prompt, mode, cfg_scale=7.0, steps=30, seed=None, width=512, height=768, lora_scale=0.85):
     try:
         with torch_gc():
             # 한글 처리
             if contains_korean(prompt):
-                translator = get_translator()
+                translator = model_manager.get_translator()
                 with torch.inference_mode():
                     translated = translator(prompt)[0]['translation_text']
                     actual_prompt = translated
             else:
                 actual_prompt = prompt
 
+
             # 파이프라인 초기화
             pipe = DiffusionPipeline.from_pretrained(
                 BASE_MODEL,
@@ -375,8 +413,10 @@ def generate_image(prompt, mode, cfg_scale=7.0, steps=30, seed=None, width=512,
             del pipe
             return result, seed
 
+
     except Exception as e:
         raise gr.Error(f"Generation failed: {str(e)}")
+
         
 # 초기 설정 실행
 setup()