Add hiresfix

app.py

@@ -20,6 +20,8 @@ SYSTEM_PROMPT = r"""This two-panel split-frame image showcases a furniture in as
 [LEFT] standalone product shot image the furniture on a white background.
 [RIGHT] integrated example within a room scene."""
 
+MASK_CONTEXT_PADDING = 16 * 8
+
 if not torch.cuda.is_available():
 
     def _dummy_pipe(image: Image.Image, *args, **kwargs):  # noqa: ARG001
@@ -78,68 +80,135 @@ def make_example(image_path: Path, mask_path: Path) -> EditorValue:
     }
 
 
+def remove_padding(image, original_size):
+    # Get current dimensions
+    padded_width, padded_height = image.size
+    original_width, original_height = original_size
+
+    # Calculate cropping box
+    left = (padded_width - original_width) // 2
+    top = (padded_height - original_height) // 2
+    right = left + original_width
+    bottom = top + original_height
+
+    # Crop to original size
+    return image.crop((left, top, right, bottom))
+
+
 @spaces.GPU(duration=150)
 def infer(
-    furniture_image: Image.Image,
-    room_image: EditorValue,
-    prompt: str = "",
+    furniture_image_input: Image.Image,
+    room_image_input: EditorValue,
+    furniture_prompt: str = "",
     seed: int = 42,
     randomize_seed: bool = False,
     guidance_scale: float = 3.5,
     num_inference_steps: int = 20,
     max_dimension: int = 720,
+    num_images_per_prompt: int = 2,
     progress: gr.Progress = gr.Progress(track_tqdm=True),  # noqa: ARG001, B008
 ):
     # Ensure max_dimension is a multiple of 16 (for VAE)
     max_dimension = (max_dimension // 16) * 16
 
-    _room_image = room_image["background"]
-    if _room_image is None:
+    room_image = room_image_input["background"]
+    if room_image is None:
         msg = "Room image is required"
         raise ValueError(msg)
-    _room_image = cast("Image.Image", _room_image)
-    _room_image = ImageOps.fit(
-        _room_image,
+    room_image = cast("Image.Image", room_image)
+
+    room_mask = room_image_input["layers"][0]
+    if room_mask is None:
+        msg = "Room mask is required"
+        raise ValueError(msg)
+    room_mask = cast("Image.Image", room_mask)
+
+    mask_bbox_x_min, mask_bbox_y_min, mask_bbox_x_max, mask_bbox_y_max = (
+        room_mask.getbbox(alpha_only=False)
+    )
+    # Add MASK_CONTEXT_PADDING (16 pixels) for the context
+    mask_bbox_x_min -= MASK_CONTEXT_PADDING
+    mask_bbox_x_min = max(mask_bbox_x_min, 0)
+    mask_bbox_y_min -= MASK_CONTEXT_PADDING
+    mask_bbox_y_min = max(mask_bbox_y_min, 0)
+    mask_bbox_x_max += MASK_CONTEXT_PADDING
+    mask_bbox_x_max = min(mask_bbox_x_max, room_mask.width)
+    mask_bbox_y_max += MASK_CONTEXT_PADDING
+    mask_bbox_y_max = min(mask_bbox_y_max, room_mask.height)
+
+    bbox_longest_side = max(
+        mask_bbox_x_max - mask_bbox_x_min,
+        mask_bbox_y_max - mask_bbox_y_min,
+    )
+
+    room_image_cropped = room_image.crop((
+        mask_bbox_x_min,
+        mask_bbox_y_min,
+        mask_bbox_x_max,
+        mask_bbox_y_max,
+    ))
+    room_image_cropped = ImageOps.pad(
+        room_image_cropped,
+        (bbox_longest_side, bbox_longest_side),
+        # White padding
+        color=(255, 255, 255),
+        centering=(0.5, 0.5),
+    )
+    room_image_cropped = ImageOps.fit(
+        room_image_cropped,
         (max_dimension, max_dimension),
         method=Image.Resampling.LANCZOS,
         centering=(0.5, 0.5),
     )
 
-    _room_mask = room_image["layers"][0]
-    if _room_mask is None:
-        msg = "Room mask is required"
-        raise ValueError(msg)
-    _room_mask = cast("Image.Image", _room_mask)
-    _room_mask = ImageOps.fit(
-        _room_mask,
+    room_mask_cropped = room_mask.crop((
+        mask_bbox_x_min,
+        mask_bbox_y_min,
+        mask_bbox_x_max,
+        mask_bbox_y_max,
+    ))
+    # room_mask_cropped.save("room_mask_croppedv1.png")
+    room_mask_cropped = ImageOps.pad(
+        room_mask_cropped,
+        (max_dimension, max_dimension),
+        # White padding
+        color=(255, 255, 255),
+        centering=(0.5, 0.5),
+    )
+    room_mask_cropped = ImageOps.fit(
+        room_mask_cropped,
         (max_dimension, max_dimension),
         method=Image.Resampling.LANCZOS,
         centering=(0.5, 0.5),
     )
 
-    # _room_image.save("room_image.png")
+    # room_image_cropped.save("room_image_cropped.png")
+    # room_mask_cropped.save("room_mask_cropped.png")
+
+    # _room_image = ImageOps.fit(
+    #     _room_image,
+    #     (max_dimension, max_dimension),
+    #     method=Image.Resampling.LANCZOS,
+    #     centering=(0.5, 0.5),
+    # )
+    _room_image.save("room_image.png")
     # _room_mask_with_white_background = Image.new(
     #     "RGB", _room_mask.size, (255, 255, 255)
     # )
     # _room_mask_with_white_background.paste(_room_mask, (0, 0), _room_mask)
-    # _room_mask_with_white_background.save("room_mask.png")
+    _room_mask_with_white_background.save("room_mask.png")
 
-    furniture_image = ImageOps.fit(
-        furniture_image,
+    furniture_image = ImageOps.pad(
+        furniture_image_input,
         (max_dimension, max_dimension),
-        method=Image.Resampling.LANCZOS,
+        # White padding
+        color=(255, 255, 255),
         centering=(0.5, 0.5),
     )
-    _furniture_image = Image.new(
-        "RGB",
-        (max_dimension, max_dimension),
-        (255, 255, 255),
-    )
-    _furniture_image.paste(furniture_image, (0, 0))
 
-    # _furniture_image.save("furniture_image.png")
+    _furniture_image.save("furniture_image.png")
 
-    _furniture_mask = Image.new("RGB", (max_dimension, max_dimension), (255, 255, 255))
+    furniture_mask = Image.new("RGB", (max_dimension, max_dimension), (255, 255, 255))
 
     image = Image.new(
         "RGB",
@@ -147,16 +216,16 @@ def infer(
         (255, 255, 255),
     )
     # Paste on the center of the image
-    image.paste(_furniture_image, (0, 0))
-    image.paste(_room_image, (max_dimension, 0))
+    image.paste(furniture_image, (0, 0))
+    image.paste(room_image_cropped, (max_dimension, 0))
 
     mask = Image.new(
         "RGB",
         (max_dimension * 2, max_dimension),
         (255, 255, 255),
     )
-    mask.paste(_furniture_mask, (0, 0))
-    mask.paste(_room_mask, (max_dimension, 0), _room_mask)
+    mask.paste(furniture_mask, (0, 0))
+    mask.paste(room_mask_cropped, (max_dimension, 0), room_mask_cropped)
     # Invert the mask
     mask = ImageOps.invert(mask)
     # Blur the mask
@@ -167,7 +236,11 @@ def infer(
     if randomize_seed:
         seed = secrets.randbelow(MAX_SEED)
 
-    prompt = prompt + ".\n" + SYSTEM_PROMPT if prompt else SYSTEM_PROMPT
+    prompt = (
+        furniture_prompt + ".\n" + SYSTEM_PROMPT if furniture_prompt else SYSTEM_PROMPT
+    )
+    # image.save("image.png")
+    # mask.save("mask.png")
     results_images = pipe(
         prompt=prompt,
         image=image,
@@ -176,16 +249,36 @@ def infer(
         width=max_dimension * 2,
         num_inference_steps=num_inference_steps,
         guidance_scale=guidance_scale,
-        num_images_per_prompt=2,
+        num_images_per_prompt=num_images_per_prompt,
         generator=torch.Generator("cpu").manual_seed(seed),
     )["images"]
 
-    cropped_images = [
-        image.crop((max_dimension, 0, max_dimension * 2, max_dimension))
-        for image in results_images
-    ]
+    final_images = []
+    for image in results_images:
+        final_image = room_image.copy()
 
-    return cropped_images, seed
+        # Downscale back to the bbox_longest_side
+        image_generated = image.crop((
+            max_dimension,
+            0,
+            max_dimension * 2,
+            max_dimension,
+        ))
+        image_generated = image_generated.resize((bbox_longest_side, bbox_longest_side))
+        # Crop back to the bbox (remove the padding)
+        image_generated = remove_padding(
+            image_generated,
+            (
+                mask_bbox_x_max - mask_bbox_x_min,
+                mask_bbox_y_max - mask_bbox_y_min,
+            ),
+        )
+        # Paste the image on the room image as the crop was done
+        # on the room image
+        final_image.paste(image_generated, (mask_bbox_x_min, mask_bbox_y_min))
+        final_images.append(final_image)
+
+    return final_images, seed
 
 
 intro_markdown = r"""
@@ -241,7 +334,7 @@ with gr.Blocks(css=css) as demo:
             """,
                 max_height=50,
             )
-            furniture_image = gr.Image(
+            furniture_image_input = gr.Image(
                 label="Furniture Image",
                 type="pil",
                 sources=["upload"],
@@ -254,7 +347,7 @@ with gr.Blocks(css=css) as demo:
                     EXAMPLES_DIR / "2" / "furniture_image.png",
                 ],
                 examples_per_page=12,
-                inputs=[furniture_image],
+                inputs=[furniture_image_input],
             )
         with gr.Column(elem_id="col-mid"):
             gr.HTML(
@@ -267,7 +360,7 @@ with gr.Blocks(css=css) as demo:
             """,
                 max_height=50,
             )
-            room_image = gr.ImageEditor(
+            room_image_input = gr.ImageEditor(
                 label="Room Image - Draw mask for inpainting",
                 type="pil",
                 sources=["upload"],
@@ -288,7 +381,7 @@ with gr.Blocks(css=css) as demo:
                         EXAMPLES_DIR / "2" / "room_mask.png",
                     ),
                 ],
-                inputs=[room_image],
+                inputs=[room_image_input],
             )
         with gr.Column(elem_id="col-right"):
             gr.HTML(
@@ -309,6 +402,12 @@ with gr.Blocks(css=css) as demo:
                 height=500,
             )
             run_button = gr.Button("Run")
+
+            # Reset the results when the run button is clicked
+            run_button.click(
+                outputs=results,
+                fn=lambda: None,
+            )
             with gr.Accordion("Advanced Settings", open=False):
                 seed = gr.Slider(
                     label="Seed",
@@ -334,6 +433,14 @@ with gr.Blocks(css=css) as demo:
                         value=720,
                     )
 
+                    num_images_per_prompt = gr.Slider(
+                        label="Number of images per prompt",
+                        minimum=1,
+                        maximum=4,
+                        step=1,
+                        value=2,
+                    )
+
                     guidance_scale = gr.Slider(
                         label="Guidance Scale",
                         minimum=1,
@@ -378,21 +485,23 @@ with gr.Blocks(css=css) as demo:
                     ),
                 ],
             ],
-            inputs=[furniture_image, room_image],
+            inputs=[furniture_image_input, room_image_input],
             label=None,
         )
+
     gr.on(
-        triggers=[run_button.click, furniture_prompt.submit],
+        triggers=[run_button.click],
         fn=infer,
         inputs=[
-            furniture_image,
-            room_image,
+            furniture_image_input,
+            room_image_input,
             furniture_prompt,
             seed,
             randomize_seed,
             guidance_scale,
             num_inference_steps,
             max_dimension,
+            num_images_per_prompt,
         ],
         outputs=[results, seed],
     )