Update app.py

app.py

@@ -36,11 +36,21 @@ def safe_model_call(func):
             raise
     return wrapper
 
-# 메모리 관리 함수 수정
+
+# 메모리 관리를 위한 컨텍스트 매니저
+@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 clear_memory():
     gc.collect()
-    if torch.cuda.is_available() and torch.cuda.current_device() >= 0:
-        torch.cuda.empty_cache()
+
 
 def setup_environment():
     os.environ['PYTORCH_CUDA_ALLOC_CONF'] = 'max_split_size_mb:128'
@@ -53,52 +63,48 @@ def setup_environment():
 @spaces.GPU()
 def generate_image(prompt, mode, cfg_scale=7.0, steps=30, seed=None, width=512, height=768, lora_scale=0.85):
     try:
-        # 한글 처리
-        if contains_korean(prompt):
-            translator = get_translator()
-            translated = translator(prompt)[0]['translation_text']
-            actual_prompt = translated
-        else:
-            actual_prompt = prompt
-
-        # 파이프라인 초기화
-        pipe = DiffusionPipeline.from_pretrained(
-            BASE_MODEL,
-            torch_dtype=torch.float16,
-        )
-        pipe.to("cuda")
-
-        # LoRA 설정
-        if mode == "Generate Model":
-            pipe.load_lora_weights(MODEL_LORA_REPO)
-            trigger_word = "fashion photography, professional model"
-        else:
-            pipe.load_lora_weights(CLOTHES_LORA_REPO)
-            trigger_word = "upper clothing, fashion item"
-
-        # 생성 설정
-        generator = torch.Generator("cuda").manual_seed(seed if seed is not None else torch.randint(0, 2**32 - 1, (1,)).item())
+        with torch_gc():
+            # 한글 처리
+            if contains_korean(prompt):
+                translator = get_translator()
+                with torch.inference_mode():
+                    translated = translator(prompt)[0]['translation_text']
+                    actual_prompt = translated
+            else:
+                actual_prompt = prompt
 
-        # 이미지 생성
-        with torch.inference_mode():
-            result = pipe(
-                prompt=f"{actual_prompt} {trigger_word}",
-                num_inference_steps=steps,
-                guidance_scale=cfg_scale,
-                width=width,
-                height=height,
-                generator=generator,
-                joint_attention_kwargs={"scale": lora_scale},
-            ).images[0]
+            # 파이프라인 초기화
+            pipe = DiffusionPipeline.from_pretrained(
+                BASE_MODEL,
+                torch_dtype=torch.float16,
+            )
+            pipe = pipe.to("cuda")
 
-        # 메모리 정리
-        del pipe
-        clear_memory()
+            # LoRA 설정
+            if mode == "Generate Model":
+                pipe.load_lora_weights(MODEL_LORA_REPO)
+                trigger_word = "fashion photography, professional model"
+            else:
+                pipe.load_lora_weights(CLOTHES_LORA_REPO)
+                trigger_word = "upper clothing, fashion item"
 
-        return result, seed
+            # 이미지 생성
+            with torch.inference_mode():
+                result = pipe(
+                    prompt=f"{actual_prompt} {trigger_word}",
+                    num_inference_steps=steps,
+                    guidance_scale=cfg_scale,
+                    width=width,
+                    height=height,
+                    generator=torch.Generator("cuda").manual_seed(
+                        seed if seed is not None else torch.randint(0, 2**32 - 1, (1,)).item()
+                    ),
+                    joint_attention_kwargs={"scale": lora_scale},
+                ).images[0]
+
+            return result, seed
 
     except Exception as e:
-        clear_memory()
         raise gr.Error(f"Generation failed: {str(e)}")
 
 def contains_korean(text):
@@ -109,35 +115,17 @@ def get_translator():
     return pipeline("translation", model="Helsinki-NLP/opus-mt-ko-en", device="cuda")
 
 
-# 전역 변수 초기화
-fashion_pipe = None
-translator = None
-mask_predictor = None
-densepose_predictor = None
-vt_model = None
-pt_model = None
-vt_inference = None
-pt_inference = None
-device = None
-HF_TOKEN = None
-
 # 환경 설정 실행
 setup_environment()
 
 @spaces.GPU()
 def initialize_fashion_pipe():
-    global fashion_pipe
-    if fashion_pipe is None:
-        fashion_pipe = DiffusionPipeline.from_pretrained(
+    with torch_gc():
+        pipe = DiffusionPipeline.from_pretrained(
             BASE_MODEL,
             torch_dtype=torch.float16,
-            use_auth_token=HF_TOKEN
-        ).to("cuda")
-        try:
-            fashion_pipe.enable_xformers_memory_efficient_attention()
-        except Exception as e:
-            print(f"Warning: Could not enable memory efficient attention: {e}")
-    return fashion_pipe
+        )
+        return pipe.to("cuda")
 
 def setup():
     # Leffa 체크포���트 다운로드만 수행
@@ -145,12 +133,10 @@ def setup():
     
 @spaces.GPU()
 def get_translator():
-    global translator
-    if translator is None:
-        translator = pipeline("translation", 
-                            model="Helsinki-NLP/opus-mt-ko-en",
-                            device="cuda")
-    return translator
+    with torch_gc():
+        return pipeline("translation", 
+                       model="Helsinki-NLP/opus-mt-ko-en",
+                       device="cuda")
 
 @safe_model_call
 def get_mask_predictor():
@@ -174,17 +160,13 @@ def get_densepose_predictor():
 
 @spaces.GPU()
 def get_vt_model():
-    try:
+    with torch_gc():
         model = LeffaModel(
             pretrained_model_name_or_path="./ckpts/stable-diffusion-inpainting",
             pretrained_model="./ckpts/virtual_tryon.pth"
         )
-        model = model.half().to("cuda")
-        inference = LeffaInference(model=model)
-        return model, inference
-    except Exception as e:
-        print(f"Error in get_vt_model: {str(e)}")
-        raise
+        model = model.half()
+        return model.to("cuda"), LeffaInference(model=model)
 
 @spaces.GPU()
 def get_pt_model():
@@ -381,9 +363,9 @@ def leffa_predict_pt(src_image_path, ref_image_path):
 
 # 초기 설정 실행
 setup()    
-    
 def create_interface():
-    with gr.Blocks(theme="Yntec/HaleyCH_Theme_Orange") as demo:
+    with gr.Blocks(theme="Yntec/HaleyCH_Theme_Orange") as demo:    
+
         gr.Markdown("# 🎭 FitGen:Fashion Studio & Virtual Try-on")
         
         with gr.Tabs():
@@ -581,12 +563,7 @@ def create_interface():
                         )
                         pose_transfer_gen_button = gr.Button("Generate")
 
-        # 이벤트 핸들러
-        generate_button.click(
-            fn=generate_image,
-            inputs=[prompt, mode, cfg_scale, steps, seed, width, height, lora_scale],
-            outputs=[result, seed]
-        )
+
         
         vt_gen_button.click(
             fn=leffa_predict_vt,
@@ -600,16 +577,27 @@ def create_interface():
             outputs=[pt_gen_image]
         )
 
+   
+
+        generate_button.click(
+            fn=generate_image,
+            inputs=[prompt, mode, cfg_scale, steps, seed, width, height, lora_scale],
+            outputs=[result, seed]
+        ).success(
+            fn=lambda: gc.collect(),  # 성공 후 메모리 정리
+            inputs=None,
+            outputs=None
+        )
+
         return demo
 
 if __name__ == "__main__":
-    # 환경 설정
     setup_environment()
-    
-    # 인터페이스 생성 및 실행
     demo = create_interface()
+    demo.queue()  # 큐 활성화
     demo.launch(
         server_name="0.0.0.0",
         server_port=7860,
-        share=False
+        share=False,
+        memory_target_gb=0.5  # 메모리 제한 설정
     )