Update app.py

app.py

@@ -92,13 +92,22 @@ gd_model = GroundingDinoForObjectDetection.from_pretrained(gd_model_path, torch_
 gd_model = gd_model.to(device=device)
 assert isinstance(gd_model, GroundingDinoForObjectDetection)
 
-# FLUX 파이프라인 초기화
+# 파이프라인 초기화 부분도 수정
 pipe = FluxPipeline.from_pretrained(
     "black-forest-labs/FLUX.1-dev",
     torch_dtype=torch.float16,
     use_auth_token=HF_TOKEN
 )
-pipe.enable_attention_slicing(slice_size="auto")
+
+# 메모리 및 성능 최적화 설정
+pipe.enable_attention_slicing()
+pipe.enable_vae_slicing()
+pipe.enable_xformers_memory_efficient_attention()  # xformers가 설치되어 있다면
+
+if torch.cuda.is_available():
+    pipe = pipe.to("cuda")
+    pipe.enable_model_cpu_offload()  # 메모리 효율을 위한 CPU 오프로드
+    
 
 # LoRA 가중치 로드
 pipe.load_lora_weights(
@@ -173,18 +182,22 @@ def apply_mask(img: Image.Image, mask_img: Image.Image, defringe: bool = True) -
 
 def calculate_dimensions(aspect_ratio: str, base_size: int = 512) -> tuple[int, int]:
     """선택된 비율에 따라 이미지 크기 계산"""
-    # 기본 크기를 512로 고정하고, 비율에 따라 조정
+    # FLUX 파이프라인이 지원하는 안전한 크기 사용
     if aspect_ratio == "1:1":
         width = height = 512
     elif aspect_ratio == "16:9":
-        width, height = 512, 288  # 16:9 비율 유지하면서 높이 조정
+        width, height = 576, 320  # 16:9에 가까운 안전한 크기
     elif aspect_ratio == "9:16":
-        width, height = 288, 512  # 9:16 비율 유지하면서 너비 조정
+        width, height = 320, 576  # 9:16에 가까운 안전한 크기
     elif aspect_ratio == "4:3":
-        width, height = 512, 384  # 4:3 비율 유지
+        width, height = 512, 384  # 4:3에 가까운 안전한 크기
     else:
         width = height = 512
     
+    # 8의 배수로 조정
+    width = (width // 8) * 8
+    height = (height // 8) * 8
+    
     return width, height
 
 def generate_background(prompt: str, aspect_ratio: str) -> Image.Image:
@@ -196,40 +209,31 @@ def generate_background(prompt: str, aspect_ratio: str) -> Image.Image:
         
         with timer("Background generation"):
             try:
+                # 먼저 512x512로 생성
                 with torch.inference_mode():
                     image = pipe(
                         prompt=prompt,
-                        width=width,
-                        height=height,
+                        width=512,
+                        height=512,
                         num_inference_steps=8,
                         guidance_scale=4.0,
                     ).images[0]
+                
+                # 원하는 크기로 리사이즈
+                if width != 512 or height != 512:
+                    image = image.resize((width, height), Image.LANCZOS)
                 return image
+                
             except Exception as e:
                 print(f"Pipeline error: {str(e)}")
-                # 에러 발생 시 기본 크기로 재시도
-                try:
-                    with torch.inference_mode():
-                        image = pipe(
-                            prompt=prompt,
-                            width=512,
-                            height=512,
-                            num_inference_steps=8,
-                            guidance_scale=4.0,
-                        ).images[0]
-                        
-                        # 원하는 크기로 리사이즈
-                        if width != 512 or height != 512:
-                            image = image.resize((width, height), Image.LANCZOS)
-                        return image
-                except Exception as e2:
-                    print(f"Fallback generation error: {str(e2)}")
-                    return Image.new('RGB', (width, height), 'white')
+                # 에러 발생 시 흰색 배경 반환
+                return Image.new('RGB', (width, height), 'white')
                 
     except Exception as e:
         print(f"Background generation error: {str(e)}")
         return Image.new('RGB', (512, 512), 'white')
 
+
 def adjust_size_to_multiple_of_8(width: int, height: int) -> tuple[int, int]:
     """이미지 크기를 8의 배수로 조정"""
     new_width = max(8, ((width + 7) // 8) * 8)  # 최소 8픽셀 보장