Update app.py

app.py

@@ -23,63 +23,55 @@ import gc
 from PIL import Image, ImageDraw, ImageFont
 
 def clear_memory():
-    """메모리 정리 함수 - Spaces GPU 환경에 맞게 수정"""
+    """메모리 정리 함수"""
     gc.collect()
-    if torch.cuda.is_available():
-        try:
-            with torch.cuda.device('cuda:0'):  # 명시적으로 cuda:0 지정
-                torch.cuda.empty_cache()
-        except Exception as e:
-            print(f"GPU memory management warning: {e}")
+    torch.cuda.empty_cache()
+    
+    # 사용하지 않는 캐시 정리
+    if hasattr(torch.cuda, 'empty_cache'):
+        torch.cuda.empty_cache()
+    if hasattr(torch.cuda, 'ipc_collect'):
+        torch.cuda.ipc_collect()
 
 def initialize_models():
-    """모델 초기화 함수 - Spaces GPU 환경에 맞게 수정"""
     global segmenter, gd_model, gd_processor, pipe, translator
     
     try:
-        # GPU 메모리 정리
         clear_memory()
         
-        # 번역 모델은 CPU에서만 실행
-        model = AutoModelForSeq2SeqLM.from_pretrained(model_name).to('cpu')
+        # CPU에서만 실행되는 번역 모델 
+        model = AutoModelForSeq2SeqLM.from_pretrained(model_name, 
+            low_cpu_mem_usage=True).to('cpu')
         tokenizer = AutoTokenizer.from_pretrained(model_name)
-        translator = pipeline("translation", model=model, tokenizer=tokenizer, device=-1)
+        translator = pipeline("translation", model=model, tokenizer=tokenizer, 
+            device=-1)
         
-        # GroundingDINO 모델 초기화
+        # GroundingDINO 모델
         gd_processor = GroundingDinoProcessor.from_pretrained(gd_model_path)
         gd_model = GroundingDinoForObjectDetection.from_pretrained(
             gd_model_path,
             torch_dtype=torch.float16,
-            device_map='cuda:0'  # 명시적으로 cuda:0 지정
+            device_map='auto',
+            low_cpu_mem_usage=True
         )
         
-        # Segmenter 초기화
-        segmenter = BoxSegmenter(device='cuda:0')  # 명시적으로 cuda:0 지정
+        # Segmenter
+        segmenter = BoxSegmenter(device='cpu')
         
-        # FLUX 파이프라인 초기화
+        # FLUX 파이프라인
         pipe = FluxPipeline.from_pretrained(
             "black-forest-labs/FLUX.1-dev",
             torch_dtype=torch.float16,
-            token=HF_TOKEN
+            device_map='auto',
+            low_cpu_mem_usage=True
         )
-        pipe.enable_attention_slicing(slice_size="auto")
+        pipe.enable_attention_slicing()
+        pipe.enable_model_cpu_offload()
         
-        # LoRA 가중치 로드
-        pipe.load_lora_weights(
-            hf_hub_download(
-                "ByteDance/Hyper-SD",
-                "Hyper-FLUX.1-dev-8steps-lora.safetensors",
-                token=HF_TOKEN
-            )
-        )
-        pipe.fuse_lora(lora_scale=0.125)
-        
-        if torch.cuda.is_available():
-            pipe = pipe.to('cuda:0')  # 명시적으로 cuda:0 지정
-            
     except Exception as e:
         print(f"Model initialization error: {str(e)}")
         raise
+        
 # GPU 설정
 device = torch.device("cuda:0" if torch.cuda.is_available() else "cpu")  # 명시적으로 cuda:0 지정
 
@@ -234,32 +226,29 @@ def calculate_dimensions(aspect_ratio: str, base_size: int = 512) -> tuple[int,
         return base_size * 4 // 3, base_size
     return base_size, base_size
 
-@spaces.GPU(duration=20)  # 40초에서 20초로 감소
+@spaces.GPU(duration=20)
 def generate_background(prompt: str, aspect_ratio: str) -> Image.Image:
     try:
         width, height = calculate_dimensions(aspect_ratio)
         width, height = adjust_size_to_multiple_of_8(width, height)
         
-        max_size = 768
+        # 최대 크기 제한
+        max_size = 512  # 768에서 512로 감소
         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.inference_mode():
-                    image = pipe(
-                        prompt=prompt,
-                        width=width,
-                        height=height,
-                        num_inference_steps=8,
-                        guidance_scale=4.0
-                    ).images[0]
-            except Exception as e:
-                print(f"Pipeline error: {str(e)}")
-                return Image.new('RGB', (width, height), 'white')
+        with torch.inference_mode():
+            image = pipe(
+                prompt=prompt,
+                width=width,
+                height=height,
+                num_inference_steps=4,  # 8에서 4로 감소
+                guidance_scale=4.0,
+                batch_size=1
+            ).images[0]
                 
         return image
     except Exception as e:
@@ -460,7 +449,7 @@ def update_box_button(img, box_input):
     except:
         return gr.update(interactive=False, variant="secondary")
 
-def process_image(img: Image.Image, max_size: int = 1024) -> Image.Image:
+def process_image(img: Image.Image, max_size: int = 768) -> Image.Image:
     """이미지 크기 최적화"""
     if img.width > max_size or img.height > max_size:
         ratio = max_size / max(img.width, img.height)
@@ -824,14 +813,15 @@ if __name__ == "__main__":
             queue=True
         )
 
-    demo.queue(max_size=3)
-    demo.launch(
-        server_name="0.0.0.0",
-        server_port=7860,
-        share=False,
-        max_threads=2,
-        enable_queue=True,
-        cache_examples=False,
-        show_error=True,
-        show_tips=False
-    )
+demo.launch(
+    server_name="0.0.0.0",
+    server_port=7860,
+    share=False,
+    max_threads=2,
+    enable_queue=True,
+    cache_examples=False,
+    show_error=True,
+    show_tips=False,
+    max_size=1,  # 큐 크기 제한
+    memory_limit="48Gi"  # 메모리 제한 설정
+)