Update app.py

app.py

@@ -123,29 +123,39 @@ def adjust_size_to_multiple_of_8(width: int, height: int) -> tuple[int, int]:
     new_height = ((height + 7) // 8) * 8
     return new_width, new_height
 
-def generate_background(prompt: str, width: int, height: int) -> Image.Image:
+def calculate_dimensions(aspect_ratio: str, base_size: int = 512) -> tuple[int, int]:
+    """선택된 비율에 따라 이미지 크기 계산"""
+    if aspect_ratio == "1:1":
+        return base_size, base_size
+    elif aspect_ratio == "16:9":
+        return base_size * 16 // 9, base_size
+    elif aspect_ratio == "9:16":
+        return base_size, base_size * 16 // 9
+    elif aspect_ratio == "4:3":
+        return base_size * 4 // 3, base_size
+    return base_size, base_size
+
+def generate_background(prompt: str, aspect_ratio: str) -> Image.Image:
     """배경 이미지 생성 함수"""
     try:
-        # 이미지 크기를 8의 배수로 조정
-        adjusted_width, adjusted_height = adjust_size_to_multiple_of_8(width, height)
+        # 선택된 비율에 따라 크기 계산
+        width, height = calculate_dimensions(aspect_ratio)
+        
+        # 8의 배수로 조정
+        width, height = adjust_size_to_multiple_of_8(width, height)
         
         with timer("Background generation"):
             image = pipe(
                 prompt=prompt,
-                width=adjusted_width,
-                height=adjusted_height,
+                width=width,
+                height=height,
                 num_inference_steps=8,
                 guidance_scale=4.0,
             ).images[0]
-            
-            # 원본 크기로 리사이즈
-            if adjusted_width != width or adjusted_height != height:
-                image = image.resize((width, height), Image.Resampling.LANCZOS)
                 
         return image
     except Exception as e:
         raise gr.Error(f"Background generation failed: {str(e)}")
-        
 
 
 def combine_with_background(foreground: Image.Image, background: Image.Image) -> Image.Image:
@@ -170,7 +180,7 @@ def _gpu_process(img: Image.Image, prompt: str | BoundingBox | None) -> tuple[Im
     time_log.append(f"segment: {time.time() - t0}")
     return mask, bbox, time_log
 
-def _process(img: Image.Image, prompt: str | BoundingBox | None, bg_prompt: str | None = None) -> tuple[tuple[Image.Image, Image.Image, Image.Image], gr.DownloadButton]:
+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)
@@ -183,7 +193,7 @@ def _process(img: Image.Image, prompt: str | BoundingBox | None, bg_prompt: str
         masked_alpha = apply_mask(img, mask, defringe=True)
 
         if bg_prompt:
-            background = generate_background(bg_prompt, img.width, img.height)
+            background = generate_background(bg_prompt, aspect_ratio)
             combined = combine_with_background(masked_alpha, background)
         else:
             combined = Image.alpha_composite(Image.new("RGBA", masked_alpha.size, "white"), masked_alpha)
@@ -201,8 +211,6 @@ def _process(img: Image.Image, prompt: str | BoundingBox | None, bg_prompt: str
     except Exception as e:
         raise gr.Error(f"Processing failed: {str(e)}")
 
-
-
 def on_change_bbox(prompts: dict[str, Any] | None):
     return gr.update(interactive=prompts is not None)
 
@@ -210,13 +218,13 @@ def on_change_bbox(prompts: dict[str, Any] | None):
 def on_change_prompt(img: Image.Image | None, prompt: str | None, bg_prompt: str | None = None):
     return gr.update(interactive=bool(img and prompt))
 
-def process_prompt(img: Image.Image, prompt: str, bg_prompt: str | None = None) -> tuple[Image.Image, Image.Image]:
+def process_prompt(img: Image.Image, prompt: str, bg_prompt: str | None = None, aspect_ratio: str = "1:1") -> tuple[Image.Image, Image.Image]:
     try:
         if img is None or prompt.strip() == "":
             raise gr.Error("Please provide both image and prompt")
         
         # Process the image
-        results, _ = _process(img, prompt, bg_prompt)
+        results, _ = _process(img, prompt, bg_prompt, aspect_ratio)
         
         # 합성된 이미지와 추출된 이미지만 반환
         return results[1], results[2]
@@ -336,11 +344,21 @@ with gr.Blocks(theme=gr.themes.Soft(), css=css) as demo:
                 placeholder="Enter what you want to extract...",
                 interactive=True
             )
-            bg_prompt = gr.Textbox(
-                label="Background Prompt (optional)",
-                placeholder="Describe the background...",
-                interactive=True
-            )
+            with gr.Row():
+                bg_prompt = gr.Textbox(
+                    label="Background Prompt (optional)",
+                    placeholder="Describe the background...",
+                    interactive=True,
+                    scale=3
+                )
+                aspect_ratio = gr.Dropdown(
+                    choices=["1:1", "16:9", "9:16", "4:3"],
+                    value="1:1",
+                    label="Aspect Ratio",
+                    interactive=True,
+                    visible=True,
+                    scale=1
+                )
             process_btn = gr.Button(
                 "Process",
                 variant="primary",
@@ -378,12 +396,23 @@ with gr.Blocks(theme=gr.themes.Soft(), css=css) as demo:
         queue=False
     )
 
+    # bg_prompt가 비어있을 때 aspect_ratio를 비활성화하는 함수
+    def update_aspect_ratio(bg_prompt):
+        return gr.update(visible=bool(bg_prompt))
+
+    bg_prompt.change(
+        fn=update_aspect_ratio,
+        inputs=bg_prompt,
+        outputs=aspect_ratio,
+        queue=False
+    )
+
     process_btn.click(
         fn=process_prompt,
-        inputs=[input_image, text_prompt, bg_prompt],
+        inputs=[input_image, text_prompt, bg_prompt, aspect_ratio],
         outputs=[combined_image, extracted_image],
         queue=True
     )
 
 demo.queue(max_size=30, api_open=False)
-demo.launch()
+demo.launch()