Update app.py

app.py

@@ -1,7 +1,6 @@
 import os
 import numpy as np
 from typing import cast
-from pydantic import NonNegativeInt
 import torch
 from PIL import Image, ImageOps
 from diffusers import DiffusionPipeline
@@ -27,6 +26,7 @@ pipeline = DiffusionPipeline.from_pretrained(
     torch_dtype=torch.bfloat16,
     custom_pipeline=SUB_MODEL_REPO_ID,
 ).to(DEVICE)
+pipeline.post_init()
 
 
 def crop_divisible_by_16(image: Image.Image) -> Image.Image:
@@ -38,11 +38,12 @@ def crop_divisible_by_16(image: Image.Image) -> Image.Image:
 
 @spaces.GPU(duration=150)
 def predict(
-    image_and_mask: EditorValue | NonNegativeInt,
-    furniture_reference: Image.Image | None,
+    image_and_mask: EditorValue,
+    condition_image: Image.Image | None,
     seed: int = 0,
     num_inference_steps: int = 28,
-    max_dimension: int = 704,
+    condition_size: int = 512,
+    target_size: int = 512,
     condition_scale: float = 1.0,
     progress: gr.Progress = gr.Progress(track_tqdm=True),  # noqa: ARG001, B008
 ) -> Image.Image | None:
@@ -50,9 +51,15 @@ def predict(
     if not image_and_mask:
         gr.Info("Please upload an image and draw a mask")
         return None
-    if not furniture_reference:
+    if not condition_image:
         gr.Info("Please upload a furniture reference image")
         return None
+
+    pipeline.load(
+        SUB_MODEL_REPO_ID,
+        subfolder=SUB_MODEL_SUBFOLDER,
+    )
+
     image_np = image_and_mask["background"]
     image_np = cast(np.ndarray, image_np)
 
@@ -75,36 +82,45 @@ def predict(
         subfolder=SUB_MODEL_SUBFOLDER,
     )
 
-    image = Image.fromarray(image_np)
+    target_image = Image.fromarray(image_np)
     # Resize to max dimension
-    image.thumbnail((max_dimension, max_dimension))
+    target_image.thumbnail((target_size, target_size))
+    target_image = target_image.resize((target_size, target_size))
     # Ensure dimensions are multiple of 16 (for VAE)
-    image = crop_divisible_by_16(image)
+    target_image = crop_divisible_by_16(target_image)
 
-    mask = Image.fromarray(mask_np)
-    mask.thumbnail((max_dimension, max_dimension))
-    mask = crop_divisible_by_16(mask)
+    mask_image = Image.fromarray(mask_np)
+    mask_image.thumbnail((target_size, target_size))
+    mask_image = mask_image.resize((target_size, target_size))
+    mask_image = crop_divisible_by_16(mask_image)
     # Invert the mask
-    mask = ImageOps.invert(mask)
+    mask_image = ImageOps.invert(mask_image)
 
-    # Image masked is the image with the mask applied (black background)
-    image_masked = Image.new("RGB", image.size, (0, 0, 0))
-    image_masked.paste(image, (0, 0), mask)
+    # # Image masked is the image with the mask applied (black background)
+    # image_masked = Image.new("RGB", image.size, (0, 0, 0))
+    # image_masked.paste(image, (0, 0), mask)
 
-    furniture_reference.thumbnail((max_dimension, max_dimension))
-    furniture_reference = crop_divisible_by_16(furniture_reference)
+    condition_image.thumbnail((condition_size, condition_size))
+    condition_image = condition_image.resize((condition_size, condition_size))
+    condition_image = crop_divisible_by_16(condition_image)
 
     generator = torch.Generator(device="cpu").manual_seed(seed)
 
     final_image = pipeline(
-        condition_image=image_masked,
-        reference_image=furniture_reference,
-        condition_scale=condition_scale,
+        condition_image=condition_image,
         prompt="",
+        image=target_image,
+        mask_image=mask_image,
         num_inference_steps=num_inference_steps,
+        height=target_size,
+        width=target_size,
+        union_cond_attn=True,
+        add_cond_attn=False,
+        latent_lora=False,
+        default_lora=False,
+        condition_scale=condition_scale,
         generator=generator,
         max_sequence_length=512,
-        latent_lora=True,
     ).images[0]
 
     return final_image
@@ -157,7 +173,7 @@ with gr.Blocks(css=css) as demo:
                 brush=gr.Brush(default_size=75, colors=["#000000"], color_mode="fixed"),
                 transforms=[],
             )
-            furniture_reference = gr.Image(
+            condition_image = gr.Image(
                 label="Furniture Reference",
                 type="pil",
                 sources=["upload"],
@@ -197,12 +213,20 @@ with gr.Blocks(css=css) as demo:
                     value=1.0,
                 )
                 with gr.Column():
-                    max_dimension = gr.Slider(
-                        label="Max Dimension",
-                        minimum=512,
-                        maximum=2048,
+                    condition_size = gr.Slider(
+                        label="Condition Size",
+                        minimum=256,
+                        maximum=1024,
+                        step=128,
+                        value=512,
+                    )
+
+                    target_size = gr.Slider(
+                        label="Target Size",
+                        minimum=256,
+                        maximum=1024,
                         step=128,
-                        value=704,
+                        value=512,
                     )
 
                     num_inference_steps = gr.Slider(
@@ -217,10 +241,11 @@ with gr.Blocks(css=css) as demo:
         fn=predict,
         inputs=[
             image_and_mask,
-            furniture_reference,
+            condition_image,
             seed,
             num_inference_steps,
-            max_dimension,
+            condition_size,
+            target_size,
             condition_scale,
         ],
         # outputs=[image_slider],