Update app.py

app.py

@@ -18,6 +18,32 @@ import gc
 # 메모리 관리 설정 추가
 import torch.backends.cuda
 torch.backends.cuda.max_split_size_mb = 128  # 메모리 분할 크기 제한
+
+
+# 전역 변수로 모델들을 선언
+fashion_pipe = None
+translator = None
+mask_predictor = None
+densepose_predictor = None
+vt_model = None
+pt_model = None
+vt_inference = None
+pt_inference = None
+
+# 초기화 함수
+def initialize_models():
+    global fashion_pipe
+    if fashion_pipe is None:
+        fashion_pipe = DiffusionPipeline.from_pretrained(
+            BASE_MODEL,
+            torch_dtype=torch.float16,
+            use_auth_token=HF_TOKEN
+        )
+        fashion_pipe.to(device)
+
+# 앱 시작 시 모델 초기화
+initialize_models()
+
 # 메모리 관리 설정
 torch.cuda.empty_cache()
 gc.collect()
@@ -30,6 +56,16 @@ def clear_memory():
         torch.cuda.synchronize()
     gc.collect()
 
+# 모델 사용 후 메모리 해제
+def unload_models():
+    global fashion_pipe, translator, mask_predictor, densepose_predictor, vt_model, pt_model
+    fashion_pipe = None
+    translator = None
+    mask_predictor = None
+    densepose_predictor = None
+    vt_model = None
+    pt_model = None
+    clear_memory()
 # 상수 정의
 MAX_SEED = 2**32 - 1
 BASE_MODEL = "black-forest-labs/FLUX.1-dev"
@@ -194,13 +230,98 @@ def generate_fashion(prompt, mode, cfg_scale, steps, randomize_seed, seed, width
         clear_memory()  # 오류 발생 시에도 메모리 정리
         raise e
 
+
+
+def leffa_predict(src_image_path, ref_image_path, control_type):
+    global mask_predictor, densepose_predictor, vt_model, pt_model, vt_inference, pt_inference
+    
+    clear_memory()
+    
+    try:
+        # 필요한 모델 초기화
+        if control_type == "virtual_tryon" and vt_model is None:
+            vt_model = LeffaModel(
+                pretrained_model_name_or_path="./ckpts/stable-diffusion-inpainting",
+                pretrained_model="./ckpts/virtual_tryon.pth"
+            )
+            vt_model.to(device)
+            vt_inference = LeffaInference(model=vt_model)
+        
+        elif control_type == "pose_transfer" and pt_model is None:
+            pt_model = LeffaModel(
+                pretrained_model_name_or_path="./ckpts/stable-diffusion-xl-1.0-inpainting-0.1",
+                pretrained_model="./ckpts/pose_transfer.pth"
+            )
+            pt_model.to(device)
+            pt_inference = LeffaInference(model=pt_model)
+
+        if mask_predictor is None:
+            mask_predictor = AutoMasker(
+                densepose_path="./ckpts/densepose",
+                schp_path="./ckpts/schp",
+            )
+
+        if densepose_predictor is None:
+            densepose_predictor = DensePosePredictor(
+                config_path="./ckpts/densepose/densepose_rcnn_R_50_FPN_s1x.yaml",
+                weights_path="./ckpts/densepose/model_final_162be9.pkl",
+            )
+
+        # 이미지 처리
+        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)
+
+        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_predictor(src_image, "upper")["mask"]
+        else:
+            mask = Image.fromarray(np.ones_like(src_image_array) * 255)
+
+        # DensePose 예측
+        src_image_iuv_array = densepose_predictor.predict_iuv(src_image_array)
+        src_image_seg_array = densepose_predictor.predict_seg(src_image_array)
+        src_image_iuv = Image.fromarray(src_image_iuv_array)
+        src_image_seg = Image.fromarray(src_image_seg_array)
+        
+        if control_type == "virtual_tryon":
+            densepose = src_image_seg
+            inference = vt_inference
+        else:
+            densepose = src_image_iuv
+            inference = pt_inference
+
+        # Leffa 변환 및 추론
+        transform = LeffaTransform()
+        data = {
+            "src_image": [src_image],
+            "ref_image": [ref_image],
+            "mask": [mask],
+            "densepose": [densepose],
+        }
+        data = transform(data)
+        
+        output = inference(data)
+        gen_image = output["generated_image"][0]
+        
+        clear_memory()
+        return np.array(gen_image)
+        
+    except Exception as e:
+        clear_memory()
+        raise e
+
 def leffa_predict_vt(src_image_path, ref_image_path):
     return leffa_predict(src_image_path, ref_image_path, "virtual_tryon")
 
 def leffa_predict_pt(src_image_path, ref_image_path):
     return leffa_predict(src_image_path, ref_image_path, "pose_transfer")
-
-
+    
 # Gradio 인터페이스
 with gr.Blocks(theme="Yntec/HaleyCH_Theme_Orange") as demo:
     gr.Markdown("# 🎭 Fashion Studio & Virtual Try-on")