Update app.py

app.py

@@ -175,29 +175,46 @@ def contains_korean(text):
     return any(ord('가') <= ord(char) <= ord('힣') for char in text)
 
 
-# 메인 기능 함수들
+# 모델 초기화 함수 수정
+@spaces.GPU()
+def initialize_fashion_pipe():
+    try:
+        pipe = DiffusionPipeline.from_pretrained(
+            BASE_MODEL,
+            torch_dtype=torch.float16,
+            safety_checker=None,
+            requires_safety_checker=False
+        ).to("cuda")
+        pipe.enable_model_cpu_offload()
+        return pipe
+    except Exception as e:
+        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:
         # 한글 처리
         if contains_korean(prompt):
-            translator = get_translator()
-            translated = translator(prompt)[0]['translation_text']
-            actual_prompt = translated
+            with torch.inference_mode():
+                translator = get_translator()
+                translated = translator(prompt)[0]['translation_text']
+                actual_prompt = translated
         else:
             actual_prompt = prompt
 
-        # 파이프라인 가져오기
+        # 파이프라인 초기화
         pipe = initialize_fashion_pipe()
         
         # LoRA 설정
         if mode == "Generate Model":
-            pipe = load_lora(pipe, MODEL_LORA_REPO)
+            pipe.load_lora_weights(MODEL_LORA_REPO)
             trigger_word = "fashion photography, professional model"
         else:
-            pipe = load_lora(pipe, CLOTHES_LORA_REPO)
+            pipe.load_lora_weights(CLOTHES_LORA_REPO)
             trigger_word = "upper clothing, fashion item"
-        
+
         # 파라미터 제한
         width = min(width, 768)
         height = min(height, 768)
@@ -206,27 +223,32 @@ def generate_fashion(prompt, mode, cfg_scale, steps, randomize_seed, seed, width
         # 시드 설정
         if randomize_seed:
             seed = random.randint(0, MAX_SEED)
-        generator = torch.Generator(device="cuda").manual_seed(seed)
-
-        # 진행률 표시
-        progress(0, "Starting fashion generation...")
+        generator = torch.Generator("cuda").manual_seed(seed)
 
         # 이미지 생성
-        image = 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]
+        with torch.inference_mode():
+            output = pipe(
+                prompt=f"{actual_prompt} {trigger_word}",
+                num_inference_steps=steps,
+                guidance_scale=cfg_scale,
+                width=width,
+                height=height,
+                generator=generator,
+                cross_attention_kwargs={"scale": lora_scale},
+            )
+            
+        image = output.images[0]
+        
+        # 메모리 정리
+        del pipe
+        torch.cuda.empty_cache()
+        gc.collect()
 
         return image, seed
-        
+
     except Exception as e:
         print(f"Error in generate_fashion: {str(e)}")
-        raise
+        raise gr.Error(f"Generation failed: {str(e)}")
 
 @safe_model_call
 def leffa_predict(src_image_path, ref_image_path, control_type):