Update app.py

app.py

@@ -54,10 +54,18 @@ def load_model_with_optimization(model_class, *args, **kwargs):
         model = model.half()  # FP16으로 변환
     return model.to(device)
 
-# LoRA 로드 함수
 def load_lora(pipe, lora_path):
-    pipe.load_lora_weights(lora_path)
-    return pipe
+    try:
+        pipe.unload_lora_weights()  # 기존 LoRA 가중치 제거
+    except:
+        pass
+    
+    try:
+        pipe.load_lora_weights(lora_path)
+        return pipe
+    except Exception as e:
+        print(f"Warning: Failed to load LoRA weights from {lora_path}: {e}")
+        return pipe
 
 # FLUX 모델 초기화 (필요할 때만 로드)
 fashion_pipe = None
@@ -151,6 +159,14 @@ def generate_fashion(prompt, mode, cfg_scale, steps, randomize_seed, seed, width
 
         pipe = get_fashion_pipe()
         
+        # 모드에 따른 LoRA 로딩 및 트리거워드 설정
+        if mode == "Generate Model":
+            pipe = load_lora(pipe, MODEL_LORA_REPO)
+            trigger_word = "fashion photography, professional model"
+        else:
+            pipe = load_lora(pipe, CLOTHES_LORA_REPO)
+            trigger_word = "upper clothing, fashion item"
+        
         # 메모리 사용량 제한을 위한 크기 조정
         width = min(width, 768)   # 최대 크기 제한
         height = min(height, 768)  # 최대 크기 제한
@@ -178,59 +194,6 @@ 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):
-    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)
-    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":
-        mask_pred = get_mask_predictor()
-        src_image = src_image.convert("RGB")
-        mask = mask_pred(src_image, "upper")["mask"]
-    elif control_type == "pose_transfer":
-        mask = Image.fromarray(np.ones_like(src_image_array) * 255)
-
-    # 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 변환 및 추론
-    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]
-    
-    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")
 
@@ -239,7 +202,7 @@ def leffa_predict_pt(src_image_path, ref_image_path):
 
 
 # Gradio 인터페이스
-with gr.Blocks(theme=gr.themes.Default(primary_hue=gr.themes.colors.pink, secondary_hue=gr.themes.colors.red)) as demo:
+with gr.Blocks(theme="Yntec/HaleyCH_Theme_Orange") as demo:
     gr.Markdown("# 🎭 Fashion Studio & Virtual Try-on")
     
     with gr.Tabs():