Update app.py

app.py

@@ -23,12 +23,18 @@ from diffusers import FluxPipeline
 from transformers import pipeline, AutoTokenizer, AutoModelForSeq2SeqLM
 
 import gc
+import torch.cuda.amp as amp
 
+# 메모리 관리 함수 강화
 def clear_memory():
     """메모리 정리 함수"""
+    if torch.cuda.is_available():
+        torch.cuda.empty_cache()
+        torch.cuda.synchronize()
     gc.collect()
-    torch.cuda.empty_cache()
-    
+
+# 자동 혼합 정밀도(Automatic Mixed Precision) 설정
+scaler = amp.GradScaler()
 
 
 model_name = "Helsinki-NLP/opus-mt-ko-en"
@@ -79,13 +85,17 @@ gd_model = GroundingDinoForObjectDetection.from_pretrained(gd_model_path, torch_
 gd_model = gd_model.to(device=device)
 assert isinstance(gd_model, GroundingDinoForObjectDetection)
 
-# FLUX 파이프라인 초기화 부분 수정
+# FLUX 파이프라인 초기화
 pipe = FluxPipeline.from_pretrained(
     "black-forest-labs/FLUX.1-dev",
-    torch_dtype=torch.float16,  # 메모리 사용량 감소를 위해 float16 사용
+    torch_dtype=torch.float16,  # A100에 최적화된 float16 사용
     use_auth_token=HF_TOKEN,
-    device_map="balanced"  # 'auto' 대신 'balanced' 사용
+    device_map="balanced"
 )
+pipe.enable_attention_slicing(slice_size="auto")  # 메모리 사용량 최적화
+pipe.enable_sequential_cpu_offload()  # CPU 오프로딩 활성화
+
+
 pipe.load_lora_weights(
     hf_hub_download(
         "ByteDance/Hyper-SD",
@@ -95,7 +105,8 @@ pipe.load_lora_weights(
 )
 pipe.fuse_lora(lora_scale=0.125)
 
-
+os.environ["CUDA_VISIBLE_DEVICES"] = "0"  # 단일 GPU 사용
+os.environ["PYTORCH_CUDA_ALLOC_CONF"] = "max_split_size_mb:512"  # CUDA 메모리 할당 설정
 
 
 class timer:
@@ -171,37 +182,32 @@ def calculate_dimensions(aspect_ratio: str, base_size: int = 512) -> tuple[int,
     return base_size, base_size
 
 def generate_background(prompt: str, aspect_ratio: str) -> Image.Image:
-    """배경 이미지 생성 함수"""
     try:
-        # 선택된 비율에 따라 크기 계산
         width, height = calculate_dimensions(aspect_ratio)
-        
-        # 8의 배수로 조정
         width, height = adjust_size_to_multiple_of_8(width, height)
         
-        # 프롬프트 전처리
-        if not prompt or prompt.strip() == "":
-            prompt = "plain white background"
-        
+        # A100 메모리 제한을 고려한 최대 크기 설정
+        max_size = 768
+        if width > max_size or height > max_size:
+            ratio = max_size / max(width, height)
+            width = int(width * ratio)
+            height = int(height * ratio)
+            width, height = adjust_size_to_multiple_of_8(width, height)
+
         with timer("Background generation"):
-            try:
+            with torch.cuda.amp.autocast():  # 자동 혼합 정밀도 사용
                 image = pipe(
                     prompt=prompt,
                     width=width,
                     height=height,
                     num_inference_steps=8,
                     guidance_scale=4.0,
-                    max_length=77,  # CLIP 텍스트 인코더의 최대 길이 제한
+                    max_length=77,
                 ).images[0]
-            except Exception as e:
-                print(f"Pipeline error: {str(e)}")
-                # 오류 발생 시 기본 흰색 배경 생성
-                image = Image.new('RGB', (width, height), 'white')
                 
         return image
     except Exception as e:
         print(f"Background generation error: {str(e)}")
-        # 최후의 폴백: 기본 흰색 배경 반환
         return Image.new('RGB', (512, 512), 'white')
 
 # FLUX 파이프라인 초기화 부분 수정
@@ -296,34 +302,32 @@ def _gpu_process(img: Image.Image, prompt: str | BoundingBox | None) -> tuple[Im
 
 def _process(img: Image.Image, prompt: str | BoundingBox | None, bg_prompt: str | None = None, aspect_ratio: str = "1:1") -> tuple[tuple[Image.Image, Image.Image, Image.Image], gr.DownloadButton]:
     try:
-        if img.width > 2048 or img.height > 2048:
-            orig_res = max(img.width, img.height)
-            img.thumbnail((2048, 2048))
-            if isinstance(prompt, tuple):
-                x0, y0, x1, y1 = (int(x * 2048 / orig_res) for x in prompt)
-                prompt = (x0, y0, x1, y1)
+        # 입력 이미지 크기 제한
+        max_size = 1024
+        if img.width > max_size or img.height > max_size:
+            ratio = max_size / max(img.width, img.height)
+            new_size = (int(img.width * ratio), int(img.height * ratio))
+            img = img.resize(new_size, Image.LANCZOS)
+
+        # 메모리 사용량 모니터링
+        if torch.cuda.is_available():
+            torch.cuda.reset_peak_memory_stats()
 
-        mask, bbox, time_log = _gpu_process(img, prompt)
-        masked_alpha = apply_mask(img, mask, defringe=True)
+        with torch.cuda.amp.autocast():
+            mask, bbox, time_log = _gpu_process(img, prompt)
+            masked_alpha = apply_mask(img, mask, defringe=True)
 
         if bg_prompt:
-            # 배경만 생성하��� 합성은 하지 않음
             background = generate_background(bg_prompt, aspect_ratio)
-            combined = background  # 배경 이미지만 반환
+            combined = background
         else:
             combined = Image.alpha_composite(Image.new("RGBA", masked_alpha.size, "white"), masked_alpha)
 
-        thresholded = mask.point(lambda p: 255 if p > 10 else 0)
-        bbox = thresholded.getbbox()
-        to_dl = masked_alpha.crop(bbox)
-
-        temp = tempfile.NamedTemporaryFile(delete=False, suffix=".png")
-        to_dl.save(temp, format="PNG")
-        temp.close()
+        clear_memory()  # 중간 메모리 정리
 
         return (img, combined, masked_alpha), gr.DownloadButton(value=temp.name, interactive=True)
-        
     except Exception as e:
+        clear_memory()
         raise gr.Error(f"Processing failed: {str(e)}")
 
 def on_change_bbox(prompts: dict[str, Any] | None):
@@ -683,12 +687,14 @@ with gr.Blocks(theme=gr.themes.Soft(), css=css) as demo:
             </div>
         </div>
     """)
-demo.queue(max_size=20)  # 큐 크기 제한
+    
+demo.queue(max_size=10)  # 큐 크기 제한
 demo.launch(
     server_name="0.0.0.0",
     server_port=7860,
     share=False,
     enable_queue=True,
-    max_threads=4,  # 스레드 수 제한
-    allowed_paths=["examples"]
+    max_threads=2,  # 스레드 수 제한
+    allowed_paths=["examples"],
+    memory_limit=0.8  # 메모리 사용량 제한 (80%)
 )