Fix NoneType error, remove negative_prompt, use Blocks for layout

app.py

@@ -26,62 +26,60 @@ else:
 
 device = "cuda" if torch.cuda.is_available() else "cpu"
 
-# Load the correct pipeline
+# Load the pipeline
 pipe = FluxFillPipeline.from_pretrained(
     "black-forest-labs/FLUX.1-Fill-dev",
     torch_dtype=torch.bfloat16,
     token=hf_token,
 ).to(device)
-
-# Enable CPU offloading only if GPU is available
 if torch.cuda.is_available():
     pipe.enable_model_cpu_offload()
 
 
-def inpaint(
-    image: Image.Image, mask: Image.Image, prompt: str, negative_prompt: str = ""
-):
-    image = image.convert("RGB")
-    mask = mask.convert("RGB")
-
-    mask_np = np.array(mask)
-    mask_np = cv2.cvtColor(mask_np, cv2.COLOR_RGB2GRAY)
-    _, mask_np = cv2.threshold(mask_np, 127, 255, cv2.THRESH_BINARY)
-    mask = Image.fromarray(mask_np)
+def inpaint(image: Image.Image, mask: Image.Image, prompt: str):
+    if image is None or mask is None:
+        return None, None, "Please upload both a base image and a mask."
 
     try:
+        image = image.convert("RGB")
+        mask = mask.convert("RGB")
+
+        mask_np = np.array(mask)
+        mask_np = cv2.cvtColor(mask_np, cv2.COLOR_RGB2GRAY)
+        _, mask_np = cv2.threshold(mask_np, 127, 255, cv2.THRESH_BINARY)
+        processed_mask = Image.fromarray(mask_np)
+
         output = pipe(
             prompt=prompt + ", highly detailed, realistic, sepia tone",
-            negative_prompt=negative_prompt + ", blurry, low quality",
             image=image,
-            mask_image=mask,
+            mask_image=processed_mask,
             num_inference_steps=50,
             guidance_scale=7.5,
             max_sequence_length=512,
         ).images[0]
-        return output
+        return output, processed_mask, None
     except Exception as e:
-        print(f"Error during inpainting: {e}")
-        raise
+        return None, None, f"Error during inpainting: {str(e)}"
 
 
-iface = gr.Interface(
-    fn=inpaint,
-    inputs=[
-        gr.Image(label="Base Image", type="pil"),
-        gr.Image(
-            label="Mask (white=edit, black=keep)", type="pil"
-        ),  # Removed tool="sketch"
-        gr.Textbox(
+with gr.Blocks() as demo:
+    with gr.Row():
+        base_image = gr.Image(label="Base Image", type="pil")
+        mask_image = gr.Image(label="Mask (white=edit, black=keep)", type="pil")
+        prompt = gr.Textbox(
             label="Prompt",
             value="add a golden crescent moon on the forehead, glowing red cat eyes",
-        ),
-        gr.Textbox(label="Negative Prompt (optional)", value="distorted, unnatural"),
-    ],
-    outputs=gr.Image(label="Inpainted Output"),
-    title="Flux.1 Fill-dev Inpaint",
-    description="Upload an image and mask, enter a prompt, and get a stylized output with preserved structure and added styles (e.g., sepia tone, elements).",
-)
+        )
+        submit = gr.Button("Submit")
+    with gr.Column():
+        output_image = gr.Image(label="Inpainted Output")
+        processed_mask_display = gr.Image(label="Processed Mask")
+        error_label = gr.Markdown()
+    submit.click(
+        inpaint,
+        inputs=[base_image, mask_image, prompt],
+        outputs=[output_image, processed_mask_display, error_label],
+    )
 
 if __name__ == "__main__":
-    iface.launch()
+    demo.launch()