Add examples

app.py

@@ -1,4 +1,5 @@
 import secrets
+from pathlib import Path
 from typing import cast
 
 import gradio as gr
@@ -11,10 +12,9 @@ from PIL import Image, ImageFilter, ImageOps
 
 DEVICE = "cuda"
 
+EXAMPLES_DIR = Path(__file__).parent / "examples"
+
 MAX_SEED = np.iinfo(np.int32).max
-FIXED_DIMENSION = 900
-FIXED_DIMENSION = 720
-FIXED_DIMENSION = (FIXED_DIMENSION // 16) * 16
 
 SYSTEM_PROMPT = r"""This two-panel split-frame image showcases a furniture in as a product shot versus styled in a room.
 [LEFT] standalone product shot image the furniture on a white background.
@@ -30,6 +30,7 @@ else:
     state_dict, network_alphas = FluxFillPipeline.lora_state_dict(
         pretrained_model_name_or_path_or_dict="blanchon/FluxFillFurniture",
         weight_name="pytorch_lora_weights3.safetensors",
+        torch_dtype=torch.bfloat16,
         return_alphas=True,
     )
 
@@ -43,11 +44,42 @@ else:
     FluxFillPipeline.load_lora_into_transformer(
         state_dict=state_dict,
         network_alphas=network_alphas,
+        torch_dtype=torch.bfloat16,
         transformer=pipe.transformer,
     )
     pipe.to(DEVICE)
 
 
+def make_example(image_path: Path, mask_path: Path) -> EditorValue:
+    background_image = Image.open(image_path)
+    background_image = background_image.convert("RGB")
+    background = np.array(background_image)
+
+    mask_image = Image.open(mask_path)
+    mask_image = mask_image.convert("RGB")
+
+    mask = np.array(mask_image)
+    mask = mask[:, :, 0]
+    mask = np.where(mask == 255, 0, 255)  # noqa: PLR2004
+
+    if background.shape[0] != mask.shape[0] or background.shape[1] != mask.shape[1]:
+        msg = "Background and mask must have the same shape"
+        raise ValueError(msg)
+
+    layer = np.zeros((background.shape[0], background.shape[1], 4), dtype=np.uint8)
+    layer[:, :, 3] = mask
+
+    composite = np.zeros((background.shape[0], background.shape[1], 4), dtype=np.uint8)
+    composite[:, :, :3] = background
+    composite[:, :, 3] = np.where(mask == 255, 0, 255)  # noqa: PLR2004
+
+    return {
+        "background": background,
+        "layers": [layer],
+        "composite": composite,
+    }
+
+
 @spaces.GPU(duration=150)
 def infer(
     furniture_image: Image.Image,
@@ -57,16 +89,20 @@ def infer(
     randomize_seed: bool = False,
     guidance_scale: float = 3.5,
     num_inference_steps: int = 20,
+    max_dimension: int = 720,
     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:
         msg = "Room image is required"
         raise ValueError(msg)
-    _room_image = cast(Image.Image, _room_image)
+    _room_image = cast("Image.Image", _room_image)
     _room_image = ImageOps.fit(
         _room_image,
-        (FIXED_DIMENSION, FIXED_DIMENSION),
+        (max_dimension, max_dimension),
         method=Image.Resampling.LANCZOS,
         centering=(0.5, 0.5),
     )
@@ -75,47 +111,54 @@ def infer(
     if _room_mask is None:
         msg = "Room mask is required"
         raise ValueError(msg)
-    _room_mask = cast(Image.Image, _room_mask)
+    _room_mask = cast("Image.Image", _room_mask)
     _room_mask = ImageOps.fit(
         _room_mask,
-        (FIXED_DIMENSION, FIXED_DIMENSION),
+        (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")
+
     furniture_image = ImageOps.fit(
         furniture_image,
-        (FIXED_DIMENSION, FIXED_DIMENSION),
+        (max_dimension, max_dimension),
         method=Image.Resampling.LANCZOS,
         centering=(0.5, 0.5),
     )
     _furniture_image = Image.new(
         "RGB",
-        (FIXED_DIMENSION, FIXED_DIMENSION),
+        (max_dimension, max_dimension),
         (255, 255, 255),
     )
     _furniture_image.paste(furniture_image, (0, 0))
 
-    _furniture_mask = Image.new(
-        "RGB", (FIXED_DIMENSION, FIXED_DIMENSION), (255, 255, 255)
-    )
+    # _furniture_image.save("furniture_image.png")
+
+    _furniture_mask = Image.new("RGB", (max_dimension, max_dimension), (255, 255, 255))
 
     image = Image.new(
         "RGB",
-        (FIXED_DIMENSION * 2, FIXED_DIMENSION),
+        (max_dimension * 2, max_dimension),
         (255, 255, 255),
     )
     # Paste on the center of the image
     image.paste(_furniture_image, (0, 0))
-    image.paste(_room_image, (FIXED_DIMENSION, 0))
+    image.paste(_room_image, (max_dimension, 0))
 
     mask = Image.new(
         "RGB",
-        (FIXED_DIMENSION * 2, FIXED_DIMENSION),
+        (max_dimension * 2, max_dimension),
         (255, 255, 255),
     )
     mask.paste(_furniture_mask, (0, 0))
-    mask.paste(_room_mask, (FIXED_DIMENSION, 0), _room_mask)
+    mask.paste(_room_mask, (max_dimension, 0), _room_mask)
     # Invert the mask
     mask = ImageOps.invert(mask)
     # Blur the mask
@@ -131,8 +174,8 @@ def infer(
         prompt=prompt,
         image=image,
         mask_image=mask,
-        height=FIXED_DIMENSION,
-        width=FIXED_DIMENSION * 2,
+        height=max_dimension,
+        width=max_dimension * 2,
         num_inference_steps=num_inference_steps,
         guidance_scale=guidance_scale,
         num_images_per_prompt=2,
@@ -140,106 +183,196 @@ def infer(
     )["images"]
 
     cropped_images = [
-        image.crop((FIXED_DIMENSION, 0, FIXED_DIMENSION * 2, FIXED_DIMENSION))
+        image.crop((max_dimension, 0, max_dimension * 2, max_dimension))
         for image in results_images
     ]
 
     return cropped_images, seed
 
 
-intro_markdown = """
-# AnyFurnish
-
-AnyFurnish is a tool that allows you to generate furniture images using Flux.1 Fill Dev.
+intro_markdown = r"""
+<div>
+  <div>
+    <div style="display: flex; justify-content: center; align-items: center; text-align: center; font-size: 40px;">
+      <b>AnyFurnish</b>
+    </div>
+    <br>
+    <div style="display: flex; justify-content: center; align-items: center; text-align: center;">
+        <a href="https://github.com/julien-blanchon/"><img src="https://img.shields.io/static/v1?label=Github Report&message=Github&color=green"></a> &ensp;
+    </div>
+    <br>
+    <div style="display: flex; text-align: center; font-size: 14px; padding-right: 300px; padding-left: 300px;">
+        AnyFurnish is a tool that allows you to generate furniture images using Flux.1 Fill Dev. 
+        You can upload a furniture image and a room image, and the tool will generate a new image with the furniture in the room.
+    </div>
+  </div>
+</div>
 """
 
-css = """
-#col-container {
+css = r"""
+#col-left {
     margin: 0 auto;
-    max-width: 1000px;
+    max-width: 430px;
+}
+#col-mid {
+    margin: 0 auto;
+    max-width: 430px;
+}
+#col-right {
+    margin: 0 auto;
+    max-width: 430px;
+}
+#col-showcase {
+    margin: 0 auto;
+    max-width: 1100px;
 }
 """
 
-with gr.Blocks(css=css) as demo:
-    with gr.Column(elem_id="col-container"):
-        gr.Markdown(intro_markdown)
-        with gr.Row():
-            with gr.Column():
-                with gr.Column():
-                    furniture_image = gr.Image(
-                        label="Furniture Image",
-                        type="pil",
-                        sources=["upload"],
-                        image_mode="RGB",
-                        height=400,
-                    )
-                    room_image = gr.ImageEditor(
-                        label="Room Image - Draw mask for inpainting",
-                        type="pil",
-                        sources=["upload"],
-                        image_mode="RGBA",
-                        layers=False,
-                        crop_size="1:1",
-                        brush=gr.Brush(colors=["#FFFFFF"], color_mode="fixed"),
-                        height=400,
-                    )
-                prompt = gr.Text(
-                    label="Prompt",
-                    show_label=False,
-                    max_lines=1,
-                    placeholder="Enter a custom furniture description (optional)",
-                    container=False,
-                )
-                run_button = gr.Button("Run")
 
+with gr.Blocks(css=css) as demo:
+    gr.Markdown(intro_markdown)
+    with gr.Row():
+        with gr.Column(elem_id="col-left"):
+            gr.HTML(
+                """
+            <div style="display: flex; justify-content: center; align-items: center; text-align: center; font-size: 20px;">
+                <div>
+                Step 1. Upload a furniture image ⬇️
+                </div>
+            </div>
+            """,
+                max_height=50,
+            )
+            furniture_image = gr.Image(
+                label="Furniture Image",
+                type="pil",
+                sources=["upload"],
+                image_mode="RGB",
+                height=500,
+            )
+            furniture_prompt = gr.Text(
+                label="Prompt",
+                max_lines=1,
+                placeholder="Enter a custom furniture description (optional)",
+                container=False,
+            )
+        with gr.Column(elem_id="col-mid"):
+            gr.HTML(
+                """
+            <div style="display: flex; justify-content: center; align-items: center; text-align: center; font-size: 20px;">
+                <div>
+                Step 2. Upload a room image ⬇️
+                </div>
+            </div>
+            """,
+                max_height=50,
+            )
+            room_image = gr.ImageEditor(
+                label="Room Image - Draw mask for inpainting",
+                type="pil",
+                sources=["upload"],
+                image_mode="RGBA",
+                layers=False,
+                crop_size="1:1",
+                brush=gr.Brush(colors=["#FFFFFF"], color_mode="fixed"),
+                height=500,
+            )
+        with gr.Column(elem_id="col-right"):
+            gr.HTML(
+                """
+            <div style="display: flex; justify-content: center; align-items: center; text-align: center; font-size: 20px;">
+                <div>
+                Step 3. Press Run to launch
+                </div>
+            </div>
+            """,
+                max_height=50,
+            )
             results = gr.Gallery(
                 label="Results",
                 format="png",
                 show_label=False,
                 columns=2,
-                height=700,
-            )
-
-        with gr.Accordion("Advanced Settings", open=False):
-            seed = gr.Slider(
-                label="Seed",
-                minimum=0,
-                maximum=MAX_SEED,
-                step=1,
-                value=0,
+                height=500,
             )
-
-            randomize_seed = gr.Checkbox(label="Randomize seed", value=True)
-
-            with gr.Row():
-                guidance_scale = gr.Slider(
-                    label="Guidance Scale",
-                    minimum=1,
-                    maximum=30,
-                    step=0.5,
-                    # value=50,
-                    value=3.5,
-                )
-
-                num_inference_steps = gr.Slider(
-                    label="Number of inference steps",
-                    minimum=1,
-                    maximum=50,
+            run_button = gr.Button("Run")
+            with gr.Accordion("Advanced Settings", open=False):
+                seed = gr.Slider(
+                    label="Seed",
+                    minimum=0,
+                    maximum=MAX_SEED,
                     step=1,
-                    value=20,
+                    value=0,
                 )
+                randomize_seed = gr.Checkbox(label="Randomize seed", value=True)
+                with gr.Column():
+                    max_dimension = gr.Slider(
+                        label="Max Dimension",
+                        minimum=512,
+                        maximum=1024,
+                        step=128,
+                        value=720,
+                    )
+
+                    guidance_scale = gr.Slider(
+                        label="Guidance Scale",
+                        minimum=1,
+                        maximum=30,
+                        step=0.5,
+                        # value=50,  # noqa: ERA001
+                        value=30,
+                    )
+
+                    num_inference_steps = gr.Slider(
+                        label="Number of inference steps",
+                        minimum=1,
+                        maximum=50,
+                        step=1,
+                        value=20,
+                    )
 
+    with gr.Column(elem_id="col-showcase"):
+        gr.HTML("""
+        <div style="display: flex; justify-content: center; align-items: center; text-align: center; font-size: 20px;">
+            <div> </div>
+            <br>
+            <div>
+            AnyFurnish examples in pairs of furniture and room images
+            </div>
+        </div>
+        """)
+        show_case = gr.Examples(
+            examples=[
+                [
+                    EXAMPLES_DIR / "1" / "furniture_image.png",
+                    make_example(
+                        EXAMPLES_DIR / "1" / "room_image.png",
+                        EXAMPLES_DIR / "1" / "room_mask.png",
+                    ),
+                ],
+                [
+                    EXAMPLES_DIR / "2" / "furniture_image.png",
+                    make_example(
+                        EXAMPLES_DIR / "2" / "room_image.png",
+                        EXAMPLES_DIR / "2" / "room_mask.png",
+                    ),
+                ],
+            ],
+            inputs=[furniture_image, room_image],
+            label=None,
+        )
     gr.on(
-        triggers=[run_button.click, prompt.submit],
+        triggers=[run_button.click, furniture_prompt.submit],
         fn=infer,
         inputs=[
             furniture_image,
             room_image,
-            prompt,
+            furniture_prompt,
             seed,
             randomize_seed,
             guidance_scale,
             num_inference_steps,
+            max_dimension,
         ],
         outputs=[results, seed],
     )