Update app.py

app.py

@@ -196,26 +196,23 @@ def generate_background(prompt: str, aspect_ratio: str) -> Image.Image:
             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=8,
+                guidance_scale=4.0
+            ).images[0]
+            
         return image
+        
     except Exception as e:
         print(f"Background generation error: {str(e)}")
-        return Image.new('RGB', (512, 512), 'white')
+        return Image.new('RGB', (width, height), 'white')
+    finally:
+        clear_memory()
 
 def create_position_grid():
     return """
@@ -273,23 +270,20 @@ def combine_with_background(foreground: Image.Image, background: Image.Image,
     result.paste(scaled_foreground, (x, y), scaled_foreground)
     return result
 
-@spaces.GPU(duration=30)  # 120초에서 30초로 감소
+@spaces.GPU(duration=30)
 def _gpu_process(img: Image.Image, prompt: str | BoundingBox | None) -> tuple[Image.Image, BoundingBox | None, list[str]]:
-    time_log: list[str] = []
     try:
-        if isinstance(prompt, str):
-            t0 = time.time()
-            bbox = gd_detect(img, prompt)
-            time_log.append(f"detect: {time.time() - t0}")
-            if not bbox:
-                print(time_log[0])
-                raise gr.Error("No object detected")
-        else:
-            bbox = prompt
-        t0 = time.time()
-        mask = segmenter(img, bbox)
-        time_log.append(f"segment: {time.time() - t0}")
-        return mask, bbox, time_log
+        with torch.inference_mode():
+            if isinstance(prompt, str):
+                bbox = gd_detect(img, prompt)
+                if not bbox:
+                    raise gr.Error("No object detected in image")
+            else:
+                bbox = prompt
+                
+            mask = segmenter(img, bbox)
+            return mask, bbox, []
+            
     except Exception as e:
         print(f"GPU process error: {str(e)}")
         raise
@@ -345,37 +339,40 @@ def process_prompt(img: Image.Image, prompt: str, bg_prompt: str | None = None,
                   aspect_ratio: str = "1:1", position: str = "bottom-center", 
                   scale_percent: float = 100) -> tuple[Image.Image, Image.Image]:
     try:
-        if img is None or prompt.strip() == "":
+        if img is None or not prompt or prompt.isspace():
             raise gr.Error("Please provide both image and prompt")
         
-        print(f"Processing with position: {position}, scale: {scale_percent}")
-        
-        try:
-            prompt = translate_to_english(prompt)
-            if bg_prompt:
-                bg_prompt = translate_to_english(bg_prompt)
-        except Exception as e:
-            print(f"Translation error (continuing with original text): {str(e)}")
+        # 입력 이미지 크기 제한
+        max_size = 1024
+        if img.width > max_size or img.height > max_size:
+            ratio = max_size / max(img.width, img.height)
+            img = img.resize((int(img.width * ratio), int(img.height * ratio)), Image.LANCZOS)
         
-        results, _ = _process(img, prompt, bg_prompt, aspect_ratio)
+        # 번역 처리
+        translated_prompt = translate_to_english(prompt)
+        translated_bg_prompt = translate_to_english(bg_prompt) if bg_prompt else None
         
-        if bg_prompt:
-            try:
-                combined = combine_with_background(
-                    foreground=results[2],
-                    background=results[1],
-                    position=position,
-                    scale_percent=scale_percent
-                )
-                print(f"Combined image created with position: {position}")
-                return combined, results[2]
-            except Exception as e:
-                print(f"Combination error: {str(e)}")
-                return results[1], results[2]
-        
-        return results[1], results[2]
+        # 이미지 처리
+        with torch.inference_mode():
+            results, _ = _process(img, translated_prompt, translated_bg_prompt, aspect_ratio)
+            
+            if translated_bg_prompt:
+                try:
+                    combined = combine_with_background(
+                        foreground=results[2],
+                        background=results[1],
+                        position=position,
+                        scale_percent=scale_slider
+                    )
+                    return combined, results[2]
+                except Exception as e:
+                    print(f"Background combination error: {e}")
+                    return results[1], results[2]
+            
+            return results[1], results[2]
+            
     except Exception as e:
-        print(f"Error in process_prompt: {str(e)}")
+        print(f"Process error: {str(e)}")
         raise gr.Error(str(e))
     finally:
         clear_memory()
@@ -482,6 +479,18 @@ button.primary:hover {
 }
 """
 
+###--------------ZERO GPU 필수/ 메모리 관리 공통 --------------------###
+def clear_memory():
+    gc.collect()
+    if torch.cuda.is_available():
+        try:
+            torch.cuda.empty_cache()
+            torch.cuda.synchronize()
+        except:
+            pass
+
+
+
 # UI 구성
 # UI 구성 부분에서 process_btn을 위로 이동하고 position_grid.click 부분 제거
 
@@ -627,7 +636,7 @@ with gr.Blocks(theme=gr.themes.Soft(), css=css) as demo:
     )
 
 
-demo.queue(max_size=10)  # 큐 크기 제한
+demo.queue(max_size=5)  # 큐 크기 제한
 demo.launch(
     server_name="0.0.0.0",
     server_port=7860,