Update app.py

app.py

@@ -16,30 +16,54 @@ import os
 import random
 import gc
 
-# 상수 정의 (맨 앞으로 이동)
+# 상수 정의
 MAX_SEED = 2**32 - 1
 BASE_MODEL = "black-forest-labs/FLUX.1-dev"
 MODEL_LORA_REPO = "Motas/Flux_Fashion_Photography_Style"
 CLOTHES_LORA_REPO = "prithivMLmods/Canopus-Clothing-Flux-LoRA"
 
-# 메모리 관리 설정
-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
-torch.backends.cuda.max_split_size_mb = 128
-
-# 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 safe_model_call(func):
+    def wrapper(*args, **kwargs):
+        try:
+            clear_memory()
+            result = func(*args, **kwargs)
+            clear_memory()
+            return result
+        except Exception as e:
+            clear_memory()
+            print(f"Error in {func.__name__}: {str(e)}")
+            raise
+    return wrapper
 
-# CUDA 설정
-device = "cuda" if torch.cuda.is_available() else "cpu"
+# 메모리 관리 함수
+def clear_memory():
+    if torch.cuda.is_available():
+        torch.cuda.empty_cache()
+        torch.cuda.synchronize()
+    gc.collect()
 
-# 전역 변수로 모델들을 선언
+def setup_environment():
+    # 메모리 관리 설정
+    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
+    torch.backends.cuda.max_split_size_mb = 128
+
+    # Hugging Face 토큰 설정
+    global HF_TOKEN
+    HF_TOKEN = os.getenv("HF_TOKEN")
+    if HF_TOKEN is None:
+        raise ValueError("Please set the HF_TOKEN environment variable")
+    login(token=HF_TOKEN)
+
+    # CUDA 설정
+    global device
+    device = "cuda" if torch.cuda.is_available() else "cpu"
+
+# 전역 변수 초기화
 fashion_pipe = None
 translator = None
 mask_predictor = None
@@ -48,91 +72,31 @@ vt_model = None
 pt_model = None
 vt_inference = None
 pt_inference = None
+device = None
+HF_TOKEN = None
 
-def clear_memory():
-    if torch.cuda.is_available():
-        torch.cuda.empty_cache()
-        torch.cuda.synchronize()
-    gc.collect()
-
-# 초기화 함수
-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 호출
-initialize_models()
-
-
-
-# 모델 사용 후 메모리 해제
-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()
-
-# 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)
-
-# CUDA 설정
-device = "cuda" if torch.cuda.is_available() else "cpu"
+# 환경 설정 실행
+setup_environment()
 
-# 모델 로드 함수
-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)
-
-def load_lora(pipe, lora_path):
-    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
-def get_fashion_pipe():
+# 모델 관리 함수들
+def initialize_fashion_pipe():
     global fashion_pipe
     if fashion_pipe is None:
-        torch.cuda.empty_cache()
+        clear_memory()
         fashion_pipe = DiffusionPipeline.from_pretrained(
             BASE_MODEL,
             torch_dtype=torch.float16,
             use_auth_token=HF_TOKEN
         )
         try:
-            fashion_pipe.enable_xformers_memory_efficient_attention()  # 수정된 부분
+            fashion_pipe.enable_xformers_memory_efficient_attention()
         except Exception as e:
             print(f"Warning: Could not enable memory efficient attention: {e}")
         fashion_pipe.enable_sequential_cpu_offload()
     return fashion_pipe
 
-# 번역기 초기화 (필요할 때만 로드)
-translator = None
+@safe_model_call
 def get_translator():
     global translator
     if translator is None:
@@ -141,8 +105,7 @@ def get_translator():
                             device=device if device == "cuda" else -1)
     return translator
 
-
-# Leffa 모델 관련 함수들
+@safe_model_call
 def get_mask_predictor():
     global mask_predictor
     if mask_predictor is None:
@@ -152,6 +115,7 @@ def get_mask_predictor():
         )
     return mask_predictor
 
+@safe_model_call
 def get_densepose_predictor():
     global densepose_predictor
     if densepose_predictor is None:
@@ -161,41 +125,60 @@ def get_densepose_predictor():
         )
     return densepose_predictor
 
+@safe_model_call
 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,
+        vt_model = LeffaModel(
             pretrained_model_name_or_path="./ckpts/stable-diffusion-inpainting",
             pretrained_model="./ckpts/virtual_tryon.pth"
         )
+        vt_model = vt_model.half().to(device)
         vt_inference = LeffaInference(model=vt_model)
     return vt_model, vt_inference
 
+@safe_model_call
 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,
+        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 = pt_model.half().to(device)
         pt_inference = LeffaInference(model=pt_model)
     return pt_model, pt_inference
 
-# Leffa 체크포인트 다운로드
-snapshot_download(repo_id="franciszzj/Leffa", local_dir="./ckpts")
+def load_lora(pipe, lora_path):
+    try:
+        pipe.unload_lora_weights()
+    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
+
+# 초기 설정 함수
+def setup():
+    # Leffa 체크포인트 다운로드
+    snapshot_download(repo_id="franciszzj/Leffa", local_dir="./ckpts")
+    # 기본 모델 초기화
+    initialize_fashion_pipe()
 
+# 유틸리티 함수
 def contains_korean(text):
     return any(ord('가') <= ord(char) <= ord('힣') for char in text)
 
+
+# 메인 기능 함수들
 @spaces.GPU()
+@safe_model_call
 def generate_fashion(prompt, mode, cfg_scale, steps, randomize_seed, seed, width, height, lora_scale, progress=gr.Progress(track_tqdm=True)):
-    clear_memory()  # 메모리 정리
-    
     try:
+        # 한글 처리
         if contains_korean(prompt):
             translator = get_translator()
             translated = translator(prompt)[0]['translation_text']
@@ -203,9 +186,10 @@ def generate_fashion(prompt, mode, cfg_scale, steps, randomize_seed, seed, width
         else:
             actual_prompt = prompt
 
-        pipe = get_fashion_pipe()
+        # 파이프라인 가져오기
+        pipe = initialize_fashion_pipe()
         
-        # 모드에 따른 LoRA 로딩 및 트리거워드 설정
+        # LoRA 설정
         if mode == "Generate Model":
             pipe = load_lora(pipe, MODEL_LORA_REPO)
             trigger_word = "fashion photography, professional model"
@@ -213,19 +197,23 @@ def generate_fashion(prompt, mode, cfg_scale, steps, randomize_seed, seed, width
             pipe = load_lora(pipe, CLOTHES_LORA_REPO)
             trigger_word = "upper clothing, fashion item"
         
-        # 메모리 사용량 제한을 위한 크기 조정
-        width = min(width, 768)   # 최대 크기 제한
-        height = min(height, 768)  # 최대 크기 제한
+        # 파라미터 제한
+        width = min(width, 768)
+        height = min(height, 768)
+        steps = min(steps, 30)
         
+        # 시드 설정
         if randomize_seed:
             seed = random.randint(0, MAX_SEED)
         generator = torch.Generator(device="cuda").manual_seed(seed)
 
+        # 진행률 표시
         progress(0, "Starting fashion generation...")
 
+        # 이미지 생성
         image = pipe(
             prompt=f"{actual_prompt} {trigger_word}",
-            num_inference_steps=min(steps, 30),  # 스텝 수 제한
+            num_inference_steps=steps,
             guidance_scale=cfg_scale,
             width=width,
             height=height,
@@ -233,51 +221,25 @@ def generate_fashion(prompt, mode, cfg_scale, steps, randomize_seed, seed, width
             joint_attention_kwargs={"scale": lora_scale},
         ).images[0]
 
-        clear_memory()  # 메모리 정리
         return image, seed
         
     except Exception as e:
-        clear_memory()  # 오류 발생 시에도 메모리 정리
-        raise e
-
-
+        print(f"Error in generate_fashion: {str(e)}")
+        raise
 
+@safe_model_call
 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",
-            )
-
-        # 이미지 처리
+        # 모델 초기화
+        if control_type == "virtual_tryon":
+            model, inference = get_vt_model()
+        else:
+            model, inference = get_pt_model()
+
+        mask_pred = get_mask_predictor()
+        dense_pred = get_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)
@@ -289,22 +251,18 @@ def leffa_predict(src_image_path, ref_image_path, control_type):
         # Mask 생성
         if control_type == "virtual_tryon":
             src_image = src_image.convert("RGB")
-            mask = mask_predictor(src_image, "upper")["mask"]
+            mask = mask_pred(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)
+        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 = src_image_seg
-            inference = vt_inference
+            densepose = Image.fromarray(src_image_seg_array)
         else:
-            densepose = src_image_iuv
-            inference = pt_inference
+            densepose = Image.fromarray(src_image_iuv_array)
 
         # Leffa 변환 및 추론
         transform = LeffaTransform()
@@ -317,20 +275,22 @@ def leffa_predict(src_image_path, ref_image_path, control_type):
         data = transform(data)
         
         output = inference(data)
-        gen_image = output["generated_image"][0]
-        
-        clear_memory()
-        return np.array(gen_image)
+        return np.array(output["generated_image"][0])
         
     except Exception as e:
-        clear_memory()
-        raise e
+        print(f"Error in leffa_predict: {str(e)}")
+        raise
 
+@safe_model_call
 def leffa_predict_vt(src_image_path, ref_image_path):
     return leffa_predict(src_image_path, ref_image_path, "virtual_tryon")
 
+@safe_model_call
 def leffa_predict_pt(src_image_path, ref_image_path):
     return leffa_predict(src_image_path, ref_image_path, "pose_transfer")
+
+# 초기 설정 실행
+setup()    
     
 # Gradio 인터페이스
 with gr.Blocks(theme="Yntec/HaleyCH_Theme_Orange") as demo: