Update app.py

app.py

@@ -15,9 +15,7 @@ import gradio as gr
 import os
 import random
 import gc
-from contextlib import contextmanager  # 이 줄 추가
-
-
+from contextlib import contextmanager
 
 # 상수 정의
 MAX_SEED = 2**32 - 1
@@ -39,7 +37,6 @@ def safe_model_call(func):
             raise
     return wrapper
 
-
 # 메모리 관리를 위한 컨텍스트 매니저
 @contextmanager
 def torch_gc():
@@ -54,7 +51,6 @@ def torch_gc():
 def clear_memory():
     gc.collect()
 
-
 def setup_environment():
     os.environ['PYTORCH_CUDA_ALLOC_CONF'] = 'max_split_size_mb:128'
     HF_TOKEN = os.getenv("HF_TOKEN")
@@ -63,17 +59,6 @@ def setup_environment():
     login(token=HF_TOKEN)
     return HF_TOKEN
 
-@contextmanager
-def torch_gc():
-    try:
-        yield
-    finally:
-        gc.collect()
-        if torch.cuda.is_available() and torch.cuda.current_device() >= 0:
-            with torch.cuda.device('cuda'):
-                torch.cuda.empty_cache()
-
-
 def contains_korean(text):
     return any(ord('가') <= ord(char) <= ord('힣') for char in text)
 
@@ -81,7 +66,6 @@ def contains_korean(text):
 def get_translator():
     return pipeline("translation", model="Helsinki-NLP/opus-mt-ko-en", device="cuda")
 
-
 # 환경 설정 실행
 setup_environment()
 
@@ -97,7 +81,7 @@ def initialize_fashion_pipe():
 def setup():
     # Leffa 체크포인트 다운로드만 수행
     snapshot_download(repo_id="franciszzj/Leffa", local_dir="./ckpts")
-    
+
 @spaces.GPU()
 def get_translator():
     with torch_gc():
@@ -135,20 +119,6 @@ def get_vt_model():
         model = model.half()
         return model.to("cuda"), LeffaInference(model=model)
 
-@spaces.GPU()
-def get_pt_model():
-    try:
-        model = LeffaModel(
-            pretrained_model_name_or_path="./ckpts/stable-diffusion-xl-1.0-inpainting-0.1",
-            pretrained_model="./ckpts/pose_transfer.pth"
-        )
-        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:
         pipe.unload_lora_weights()
@@ -170,11 +140,6 @@ def get_mask_predictor():
             schp_path="./ckpts/schp",
         )
     return mask_predictor
-    
-# 유틸리티 함수
-def contains_korean(text):
-    return any(ord('가') <= ord(char) <= ord('힣') for char in text)
-
 
 # 모델 초기화 함수 수정
 @spaces.GPU()
@@ -192,7 +157,6 @@ def initialize_fashion_pipe():
         print(f"Error initializing fashion pipe: {e}")
         raise
 
-# 생성 함수 수정
 @spaces.GPU()
 def generate_fashion(prompt, mode, cfg_scale, steps, randomize_seed, seed, width, height, lora_scale, progress=gr.Progress(track_tqdm=True)):
     try:
@@ -291,10 +255,7 @@ def leffa_predict(src_image_path, ref_image_path, control_type):
     try:
         with torch_gc():
             # 모델 초기화
-            if control_type == "virtual_tryon":
-                model, inference = get_vt_model()
-            else:
-                model, inference = get_pt_model()
+            model, inference = get_vt_model()
 
             # 이미지 처리
             src_image = Image.open(src_image_path)
@@ -307,21 +268,13 @@ def leffa_predict(src_image_path, ref_image_path, control_type):
 
             # 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)
+                src_image = src_image.convert("RGB")
+                mask_pred = model_manager.get_mask_predictor()
+                mask = mask_pred(src_image, "upper")["mask"]
 
                 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)
+                densepose = Image.fromarray(src_image_seg_array)
 
             # Leffa 변환 및 추론
             transform = LeffaTransform()
@@ -348,8 +301,6 @@ def leffa_predict(src_image_path, ref_image_path, control_type):
         print(f"Error in leffa_predict: {str(e)}")
         raise
 
-
-        
 @spaces.GPU()
 def leffa_predict_vt(src_image_path, ref_image_path):
     try:
@@ -358,14 +309,6 @@ def leffa_predict_vt(src_image_path, ref_image_path):
         print(f"Error in leffa_predict_vt: {str(e)}")
         raise
 
-@spaces.GPU()
-def leffa_predict_pt(src_image_path, ref_image_path):
-    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
-
 @spaces.GPU()
 def generate_image(prompt, mode, cfg_scale=7.0, steps=30, seed=None, width=512, height=768, lora_scale=0.85):
     try:
@@ -379,7 +322,6 @@ def generate_image(prompt, mode, cfg_scale=7.0, steps=30, seed=None, width=512,
             else:
                 actual_prompt = prompt
 
-
             # 파이프라인 초기화
             pipe = DiffusionPipeline.from_pretrained(
                 BASE_MODEL,
@@ -413,16 +355,14 @@ 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()    
+
 def create_interface():
     with gr.Blocks(theme="Yntec/HaleyCH_Theme_Orange") as demo:    
-
         gr.Markdown("# 🎭 FitGen:Fashion Studio & Virtual Try-on")
         
         with gr.Tabs():
@@ -570,71 +510,11 @@ def create_interface():
                         )
                         vt_gen_button = gr.Button("Try-on")
 
-            # 포즈 전송 탭
-            with gr.Tab("Pose Transfer"):
-                with gr.Row():
-                    with gr.Column():
-                        gr.Markdown("#### Person Image")
-                        pt_ref_image = gr.Image(
-                            sources=["upload"],
-                            type="filepath",
-                            label="Person Image",
-                            width=512,
-                            height=512,
-                        )
-                        gr.Examples(
-                            inputs=pt_ref_image,
-                            examples_per_page=5,
-                            examples=["a1.webp",
-                                    "a2.webp",
-                                    "a3.webp",
-                                    "a4.webp",
-                                    "a5.webp"]
-                        )
-
-                    with gr.Column():
-                        gr.Markdown("#### Target Pose Person Image")
-                        pt_src_image = gr.Image(
-                            sources=["upload"],
-                            type="filepath",
-                            label="Target Pose Person Image",
-                            width=512,
-                            height=512,
-                        )
-                        gr.Examples(
-                            inputs=pt_src_image,
-                            examples_per_page=5,
-                            examples=["d1.webp",
-                                    "d2.webp",
-                                    "d3.webp",
-                                    "d4.webp",
-                                    "d5.webp"]
-                        )
-
-                    with gr.Column():
-                        gr.Markdown("#### Generated Image")
-                        pt_gen_image = gr.Image(
-                            label="Generated Image",
-                            width=512,
-                            height=512,
-                        )
-                        pose_transfer_gen_button = gr.Button("Generate")
-
-
-        
         vt_gen_button.click(
             fn=leffa_predict_vt,
             inputs=[vt_src_image, vt_ref_image],
             outputs=[vt_gen_image]
         )
-        
-        pose_transfer_gen_button.click(
-            fn=leffa_predict_pt,
-            inputs=[pt_src_image, pt_ref_image],
-            outputs=[pt_gen_image]
-        )
-
-   
 
         generate_button.click(
             fn=generate_image,