Update

app.py

@@ -2,17 +2,18 @@ import os
 import numpy as np
 from typing import cast
 import torch
-from PIL import Image, ImageOps
+from PIL import Image
 from diffusers import DiffusionPipeline
 import gradio as gr
 from gradio.components.image_editor import EditorValue
-import spaces
 
 DEVICE = "cuda"
 
-MAIN_MODEL_REPO_ID = os.getenv("MAIN_MODEL_REPO_ID", None)
-SUB_MODEL_REPO_ID = os.getenv("SUB_MODEL_REPO_ID", None)
-SUB_MODEL_SUBFOLDER = os.getenv("SUB_MODEL_SUBFOLDER", None)
+MAIN_MODEL_REPO_ID = os.getenv(
+    "MAIN_MODEL_REPO_ID", "black-forest-labs/FLUX.1-Fill-dev"
+)
+SUB_MODEL_REPO_ID = os.getenv("SUB_MODEL_REPO_ID", "blanchon/FurnitureModel")
+SUB_MODEL_SUBFOLDER = os.getenv("SUB_MODEL_SUBFOLDER", "model/v1/ckpt/9600")
 
 if MAIN_MODEL_REPO_ID is None:
     raise ValueError("MAIN_MODEL_REPO_ID is not set")
@@ -36,7 +37,7 @@ def crop_divisible_by_16(image: Image.Image) -> Image.Image:
     return image.crop((0, 0, w, h))
 
 
-@spaces.GPU(duration=150)
+# @spaces.GPU(duration=150)
 def predict(
     image_and_mask: EditorValue,
     condition_image: Image.Image | None,
@@ -82,35 +83,43 @@ def predict(
         subfolder=SUB_MODEL_SUBFOLDER,
     )
 
-    target_image = Image.fromarray(image_np)
+    target_image = Image.fromarray(image_np).convert("RGB")
     # Resize to 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)
-    target_image = crop_divisible_by_16(target_image)
+    new_target_image = Image.new("RGB", (target_size, target_size), (0, 0, 0))
+    new_target_image.paste(target_image, (0, 0))
 
-    mask_image = Image.fromarray(mask_np)
+    # Save target image
+    new_target_image.save("target_image.png")
+
+    mask_image = Image.fromarray(mask_np).convert("L")
     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_image = ImageOps.invert(mask_image)
+    new_mask_image = Image.new("L", (target_size, target_size), 0)
+    new_mask_image.paste(mask_image, (0, 0))
+
+    # Save mask image
+    new_mask_image.save("mask_image.png")
 
     # # 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)
 
+    condition_image = condition_image.convert("RGB")
     condition_image.thumbnail((condition_size, condition_size))
-    condition_image = condition_image.resize((condition_size, condition_size))
-    condition_image = crop_divisible_by_16(condition_image)
+    # Save condition image
+    new_condition_image = Image.new("RGB", (condition_size, condition_size), (0, 0, 0))
+    new_condition_image.paste(condition_image, (0, 0))
+
+    # Save condition image
+    new_condition_image.save("condition_image.png")
 
     generator = torch.Generator(device="cpu").manual_seed(seed)
 
     final_image = pipeline(
-        condition_image=condition_image,
+        condition_image=new_condition_image,
         prompt="",
-        image=target_image,
-        mask_image=mask_image,
+        image=new_target_image,
+        mask_image=new_mask_image,
         num_inference_steps=num_inference_steps,
         height=target_size,
         width=target_size,
@@ -123,7 +132,9 @@ def predict(
         max_sequence_length=512,
     ).images[0]
 
-    return final_image
+    final_image_crop = final_image.crop((0, 0, target_size, target_size))
+
+    return final_image_crop
 
 
 intro_markdown = r"""