Update app.py

app.py

@@ -15,13 +15,20 @@ import gradio as gr
 import os
 import random
 import gc
-
+# 메모리 관리 설정 추가
+import torch.backends.cuda
+torch.backends.cuda.max_split_size_mb = 128  # 메모리 분할 크기 제한
 # 메모리 관리 설정
 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
+def clear_memory():
+    if torch.cuda.is_available():
+        torch.cuda.empty_cache()
+        torch.cuda.synchronize()
+    gc.collect()
 
 # 상수 정의
 MAX_SEED = 2**32 - 1
@@ -63,7 +70,10 @@ def get_fashion_pipe():
             torch_dtype=torch.float16,
             use_auth_token=HF_TOKEN
         )
-        fashion_pipe.enable_memory_efficient_attention()
+        try:
+            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
 
@@ -129,51 +139,44 @@ def contains_korean(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()
+    clear_memory()  # 메모리 정리
     
-    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":
-        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"
-
-    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 = 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]
-
-    progress(100, "Completed!")
-    return image, seed
-
+    try:
+        if contains_korean(prompt):
+            translator = get_translator()
+            translated = translator(prompt)[0]['translation_text']
+            actual_prompt = translated
+        else:
+            actual_prompt = prompt
+
+        pipe = get_fashion_pipe()
+        
+        # 메모리 사용량 제한을 위한 크기 조정
+        width = min(width, 768)   # 최대 크기 제한
+        height = min(height, 768)  # 최대 크기 제한
+        
+        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),  # 스텝 수 제한
+            guidance_scale=cfg_scale,
+            width=width,
+            height=height,
+            generator=generator,
+            joint_attention_kwargs={"scale": lora_scale},
+        ).images[0]
+
+        clear_memory()  # 메모리 정리
+        return image, seed
+        
+    except Exception as e:
+        clear_memory()  # 오류 발생 시에도 메모리 정리
+        raise e
 
 def leffa_predict(src_image_path, ref_image_path, control_type):
     torch.cuda.empty_cache()