Update app.py

app.py

@@ -16,10 +16,12 @@ import os
 import random
 import gc
 
-# 메모리 최적화 설정
+# 메모리 관리 설정
+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
-os.environ['PYTORCH_CUDA_ALLOC_CONF'] = 'max_split_size_mb:512'
 
 # 상수 정의
 MAX_SEED = 2**32 - 1
@@ -27,110 +29,159 @@ BASE_MODEL = "black-forest-labs/FLUX.1-dev"
 MODEL_LORA_REPO = "Motas/Flux_Fashion_Photography_Style"
 CLOTHES_LORA_REPO = "prithivMLmods/Canopus-Clothing-Flux-LoRA"
 
-# Hugging Face 토큰 설정 및 로그인
+# Hugging Face 토큰 설정
 HF_TOKEN = os.getenv("HF_TOKEN")
 if HF_TOKEN is None:
     raise ValueError("Please set the HF_TOKEN environment variable")
 login(token=HF_TOKEN)
 
-# 메모리 정리 함수
-def clear_memory():
-    torch.cuda.empty_cache()
-    gc.collect()
-
-# 초기 메모리 정리
-clear_memory()
-
-# CUDA 사용 가능 여부 확인
+# CUDA 설정
 device = "cuda" if torch.cuda.is_available() else "cpu"
 
-# FLUX 모델 초기화
-fashion_pipe = DiffusionPipeline.from_pretrained(
-    BASE_MODEL,
-    torch_dtype=torch.float16,
-    use_auth_token=HF_TOKEN
-)
-fashion_pipe.enable_model_cpu_offload()
-
-# 번역기 초기화
-translator = pipeline("translation", model="Helsinki-NLP/opus-mt-ko-en")
+# 모델 로드 함수
+def load_model_with_optimization(model_class, *args, **kwargs):
+    torch.cuda.empty_cache()
+    gc.collect()
+    model = model_class(*args, **kwargs)
+    if device == "cuda":
+        model = model.half()  # FP16으로 변환
+    return model.to(device)
+
+# LoRA 로드 함수
+def load_lora(pipe, lora_path):
+    pipe.load_lora_weights(lora_path)
+    return pipe
+
+# FLUX 모델 초기화 (필요할 때만 로드)
+fashion_pipe = None
+def get_fashion_pipe():
+    global fashion_pipe
+    if fashion_pipe is None:
+        torch.cuda.empty_cache()
+        fashion_pipe = DiffusionPipeline.from_pretrained(
+            BASE_MODEL,
+            torch_dtype=torch.float16,
+            use_auth_token=HF_TOKEN
+        )
+        fashion_pipe.enable_memory_efficient_attention()
+        fashion_pipe.enable_sequential_cpu_offload()
+    return fashion_pipe
+
+# 번역기 초기화 (필요할 때만 로드)
+translator = None
+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)
+    return translator
+
+
+# Leffa 모델 관련 함수들
+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 get_densepose_predictor():
+    global densepose_predictor
+    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",
+        )
+    return densepose_predictor
+
+def get_vt_model():
+    global vt_model, vt_inference
+    if vt_model is None:
+        torch.cuda.empty_cache()
+        vt_model = load_model_with_optimization(
+            LeffaModel,
+            pretrained_model_name_or_path="./ckpts/stable-diffusion-inpainting",
+            pretrained_model="./ckpts/virtual_tryon.pth"
+        )
+        vt_inference = LeffaInference(model=vt_model)
+    return vt_model, vt_inference
+
+def get_pt_model():
+    global pt_model, pt_inference
+    if pt_model is None:
+        torch.cuda.empty_cache()
+        pt_model = load_model_with_optimization(
+            LeffaModel,
+            pretrained_model_name_or_path="./ckpts/stable-diffusion-xl-1.0-inpainting-0.1",
+            pretrained_model="./ckpts/pose_transfer.pth"
+        )
+        pt_inference = LeffaInference(model=pt_model)
+    return pt_model, pt_inference
 
 # Leffa 체크포인트 다운로드
 snapshot_download(repo_id="franciszzj/Leffa", local_dir="./ckpts")
 
-# Leffa 관련 모델 초기화
-mask_predictor = AutoMasker(
-    densepose_path="./ckpts/densepose",
-    schp_path="./ckpts/schp",
-)
-
-densepose_predictor = DensePosePredictor(
-    config_path="./ckpts/densepose/densepose_rcnn_R_50_FPN_s1x.yaml",
-    weights_path="./ckpts/densepose/model_final_162be9.pkl",
-)
-# Leffa 모델 초기화 수정
-vt_model = LeffaModel(
-    pretrained_model_name_or_path="./ckpts/stable-diffusion-inpainting",
-    pretrained_model="./ckpts/virtual_tryon.pth"
-)
-vt_model.to(device)  # 모델을 GPU로 이동
-vt_inference = LeffaInference(model=vt_model)
-
-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)  # 모델을 GPU로 이동
-pt_inference = LeffaInference(model=pt_model)
-
-
 def contains_korean(text):
     return any(ord('가') <= ord(char) <= ord('힣') for char in text)
 
 @spaces.GPU()
 def generate_fashion(prompt, mode, cfg_scale, steps, randomize_seed, seed, width, height, lora_scale, progress=gr.Progress(track_tqdm=True)):
+    torch.cuda.empty_cache()
+    
     if contains_korean(prompt):
+        translator = get_translator()
         translated = translator(prompt)[0]['translation_text']
         actual_prompt = translated
     else:
         actual_prompt = prompt
-        
 
     # 모드에 따른 LoRA 로딩 및 트리거워드 설정
+    pipe = get_fashion_pipe()
     if mode == "Generate Model":
-        fashion_pipe = load_lora(fashion_pipe, model_lora_repo)
+        pipe = load_lora(pipe, MODEL_LORA_REPO)
         trigger_word = "fashion photography, professional model"
     else:
-        fashion_pipe = load_lora(fashion_pipe, clothes_lora_repo)
+        pipe = load_lora(pipe, CLOTHES_LORA_REPO)
         trigger_word = "upper clothing, fashion item"
 
     if randomize_seed:
         seed = random.randint(0, MAX_SEED)
     generator = torch.Generator(device="cuda").manual_seed(seed)
 
+    # 이미지 크기 제한
+    width = min(width, 1024)
+    height = min(height, 1024)
+
     progress(0, "Starting fashion generation...")
 
     for i in range(1, steps + 1):
         if i % (steps // 10) == 0:
             progress(i / steps * 100, f"Processing step {i} of {steps}...")
 
-    image = fashion_pipe(
+    image = pipe(
         prompt=f"{actual_prompt} {trigger_word}",
         num_inference_steps=steps,
         guidance_scale=cfg_scale,
         width=width,
         height=height,
         generator=generator,
-        use_auth_token=HF_TOKEN,  # 인증 토큰 추가
         joint_attention_kwargs={"scale": lora_scale},
     ).images[0]
 
     progress(100, "Completed!")
     return image, seed
-    
+
+
 def leffa_predict(src_image_path, ref_image_path, control_type):
+    torch.cuda.empty_cache()
+    
     assert control_type in [
         "virtual_tryon", "pose_transfer"], "Invalid control type: {}".format(control_type)
+    
+    # 이미지 로드 및 크기 조정
     src_image = Image.open(src_image_path)
     ref_image = Image.open(ref_image_path)
     src_image = resize_and_center(src_image, 768, 1024)
@@ -139,26 +190,30 @@ 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
+    # Mask 생성
     if control_type == "virtual_tryon":
+        mask_pred = get_mask_predictor()
         src_image = src_image.convert("RGB")
-        mask = mask_predictor(src_image, "upper")["mask"]
+        mask = mask_pred(src_image, "upper")["mask"]
     elif control_type == "pose_transfer":
         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)
+    # DensePose 예측
+    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_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
+        model, inference = get_vt_model()
     elif control_type == "pose_transfer":
         densepose = src_image_iuv
+        model, inference = get_pt_model()
 
-    # Leffa
+    # Leffa 변환 및 추론
     transform = LeffaTransform()
-
     data = {
         "src_image": [src_image],
         "ref_image": [ref_image],
@@ -166,25 +221,21 @@ def leffa_predict(src_image_path, ref_image_path, control_type):
         "densepose": [densepose],
     }
     data = transform(data)
-    if control_type == "virtual_tryon":
-        inference = vt_inference
-    elif control_type == "pose_transfer":
-        inference = pt_inference
+    
     output = inference(data)
     gen_image = output["generated_image"][0]
-    # gen_image.save("gen_image.png")
+    
+    torch.cuda.empty_cache()
     return np.array(gen_image)
 
-
 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=gr.themes.Default(primary_hue=gr.themes.colors.pink, secondary_hue=gr.themes.colors.red)) as demo:
     gr.Markdown("# 🎭 Fashion Studio & Virtual Try-on")
     
@@ -222,7 +273,7 @@ with gr.Blocks(theme=gr.themes.Default(primary_hue=gr.themes.colors.pink, second
                                 steps = gr.Slider(
                                     label="Steps",
                                     minimum=1,
-                                    maximum=100,
+                                    maximum=50,  # 최대값 감소
                                     step=1,
                                     value=30
                                 )
@@ -238,14 +289,14 @@ with gr.Blocks(theme=gr.themes.Default(primary_hue=gr.themes.colors.pink, second
                             width = gr.Slider(
                                 label="Width",
                                 minimum=256,
-                                maximum=1536,
+                                maximum=1024,  # 최대값 감소
                                 step=64,
                                 value=512
                             )
                             height = gr.Slider(
                                 label="Height",
                                 minimum=256,
-                                maximum=1536,
+                                maximum=1024,  # 최대값 감소
                                 step=64,
                                 value=768
                             )
@@ -363,8 +414,6 @@ with gr.Blocks(theme=gr.themes.Default(primary_hue=gr.themes.colors.pink, second
                     )
                     pose_transfer_gen_button = gr.Button("Generate")
 
-    gr.Markdown(note)
-
     # 이벤트 핸들러
     generate_button.click(
         generate_fashion,
@@ -384,4 +433,5 @@ with gr.Blocks(theme=gr.themes.Default(primary_hue=gr.themes.colors.pink, second
         outputs=[pt_gen_image]
     )
 
-demo.launch(share=True, server_port=7860)
+# 앱 실행
+demo.launch(share=True, server_port=7860)