Update app.py

app.py

@@ -187,18 +187,40 @@ def remove_background_from_image(image: Image.Image) -> Image.Image:
     
     return result_image
 
+
+def calculate_dimensions(image):
+    """Calculate output dimensions based on background image, keeping largest side at 1024."""
+    if image is None:
+        return 1024, 1024
+    
+    original_width, original_height = image.size
+    
+    if original_width > original_height:
+        new_width = 1024
+        aspect_ratio = original_height / original_width
+        new_height = int(new_width * aspect_ratio)
+    else:
+        new_height = 1024
+        aspect_ratio = original_width / original_height
+        new_width = int(new_height * aspect_ratio)
+    
+    # Ensure dimensions are multiples of 8
+    new_width = (new_width // 8) * 8
+    new_height = (new_height // 8) * 8
+    
+    return new_width, new_height
+
+
 # --- Inference ---
 @spaces.GPU
 def infer(
-    gallery_images,  
+    product_image,  
     image_background,
     prompt="",
     seed=42,
     randomize_seed=True,
     true_guidance_scale=1,
     num_inference_steps=4,
-    height=None,
-    width=None,
     progress=gr.Progress(track_tqdm=True)
 ):
     if randomize_seed:
@@ -206,26 +228,27 @@ def infer(
     generator = torch.Generator(device=device).manual_seed(seed)
 
     processed_subjects = []
-    if gallery_images:
-        for img in gallery_images:
-            image = img[0]  # Extract PIL image from gallery format
+    if product_image:
+        image = remove_background_from_image(image)
             
-            image = remove_background_from_image(image)
-            
-            # Always remove alpha channels to ensure RGB format
-            image = remove_alpha_channel(image)
-            processed_subjects.append(image)
+        # Always remove alpha channels to ensure RGB format
+        image = remove_alpha_channel(image)
+        processed_subjects.append(image)
 
     all_inputs = processed_subjects
     if image_background is not None:
         all_inputs.append(image_background) 
 
+    width, height = calculate_dimensions(image_background)
+    
     if not all_inputs:
         raise gr.Error("Please upload at least one image or a background image.")
 
     result = pipe(
         image=all_inputs,
         prompt=prompt,
+        width=width,
+        height=height,
         num_inference_steps=num_inference_steps,
         generator=generator,
         true_cfg_scale=true_guidance_scale,
@@ -247,9 +270,8 @@ with gr.Blocks(theme=gr.themes.Citrus(), css=css) as demo:
         with gr.Row():
             with gr.Column():
                 with gr.Row():
-                    gallery = gr.Gallery(
-                        label="Product image (background auto removed)",
-                        columns=3, rows=2, height="auto", type="pil"
+                    product_image =  gr.Image(
+                        label="Product image (background auto removed)" type="pil"
                     )
                     image_background = gr.Image(label="Background Image", type="pil", visible=True)
                 prompt = gr.Textbox(label="Prompt")
@@ -260,8 +282,6 @@ with gr.Blocks(theme=gr.themes.Citrus(), css=css) as demo:
                     randomize_seed = gr.Checkbox(label="Randomize Seed", value=True)
                     true_guidance_scale = gr.Slider(label="True Guidance Scale", minimum=1.0, maximum=10.0, step=0.1, value=1.0)
                     num_inference_steps = gr.Slider(label="Inference Steps", minimum=1, maximum=40, step=1, value=4)
-                    height = gr.Slider(label="Height", minimum=256, maximum=2048, step=8, value=1024)
-                    width = gr.Slider(label="Width", minimum=256, maximum=2048, step=8, value=1024)
 
             with gr.Column():
                 result = gr.ImageSlider(label="Output Image", interactive=False)
@@ -274,14 +294,14 @@ with gr.Blocks(theme=gr.themes.Citrus(), css=css) as demo:
                 [[], "fusion_shoes.png", ""],
                 [["product_3.png"], "background_3.png", ""],
             ],
-            inputs=[gallery, image_background, prompt],
+            inputs=[product_image, image_background, prompt],
             outputs=[result, seed],
             fn=infer,
             cache_examples="lazy",
             elem_id="examples"
         )
 
-        inputs = [gallery, image_background, prompt, seed, randomize_seed, true_guidance_scale, num_inference_steps, height, width]
+        inputs = [product_image, image_background, prompt, seed, randomize_seed, true_guidance_scale, num_inference_steps]
         outputs = [result, seed]
 
         run_button.click(fn=infer, inputs=inputs, outputs=outputs)