Simplify interface, remove mask display and advanced settings, ensure GPU compatibility

app.py

@@ -1,85 +1,88 @@
-import os
+import gradio as gr
+import numpy as np
 import torch
 from diffusers import FluxFillPipeline
-import gradio as gr
 from PIL import Image
-import cv2
-import numpy as np
-from huggingface_hub import login
-
-# Authenticate with HF token from Spaces Secrets
-hf_token = os.getenv("HF_TOKEN")
-if not hf_token:
-    print(
-        "Warning: HF_TOKEN not found in environment. Please set it in Spaces Secrets."
-    )
-    hf_token = (
-        input("Enter your HF_TOKEN for local testing (leave blank to skip): ") or None
-    )
-
-if hf_token:
-    login(token=hf_token)
-else:
-    raise ValueError(
-        "HF_TOKEN is required to access the gated FLUX.1-Fill-dev model. Set it in Spaces Secrets or locally."
-    )
+import random
 
+# Set device (GPU if available, otherwise CPU)
 device = "cuda" if torch.cuda.is_available() else "cpu"
 
 # Load the pipeline
 pipe = FluxFillPipeline.from_pretrained(
-    "black-forest-labs/FLUX.1-Fill-dev",
-    torch_dtype=torch.bfloat16,
-    token=hf_token,
+    "black-forest-labs/FLUX.1-Fill-dev", torch_dtype=torch.bfloat16
 ).to(device)
-if torch.cuda.is_available():
-    pipe.enable_model_cpu_offload()
 
 
-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")
+# Function to calculate optimal dimensions for the output image
+def calculate_optimal_dimensions(image: Image.Image):
+    original_width, original_height = image.size
+    MIN_ASPECT_RATIO = 9 / 16
+    MAX_ASPECT_RATIO = 16 / 9
+    FIXED_DIMENSION = 1024
+    original_aspect_ratio = original_width / original_height
+    if original_aspect_ratio > 1:  # Wider than tall
+        width = FIXED_DIMENSION
+        height = round(FIXED_DIMENSION / original_aspect_ratio)
+    else:  # Taller than wide
+        height = FIXED_DIMENSION
+        width = round(FIXED_DIMENSION * original_aspect_ratio)
+    # Ensure dimensions are multiples of 8
+    width = (width // 8) * 8
+    height = (height // 8) * 8
+    # Enforce aspect ratio limits
+    calculated_aspect_ratio = width / height
+    if calculated_aspect_ratio > MAX_ASPECT_RATIO:
+        width = (height * MAX_ASPECT_RATIO // 8) * 8
+    elif calculated_aspect_ratio < MIN_ASPECT_RATIO:
+        height = (width / MIN_ASPECT_RATIO // 8) * 8
+    # Ensure width and height remain above minimum dimensions
+    width = max(width, 576) if width == FIXED_DIMENSION else width
+    height = max(height, 576) if height == FIXED_DIMENSION else height
+    return width, height
 
-        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)
 
+# Inpainting function
+def infer(edit_images, prompt):
+    image = edit_images["background"]
+    if not edit_images["layers"]:
+        raise gr.Error("Please draw a mask.")
+    mask = edit_images["layers"][0]
+    width, height = calculate_optimal_dimensions(image)
+    seed = random.randint(0, np.iinfo(np.int32).max)
+    generator = torch.Generator(device=device).manual_seed(seed)
+    try:
         output = pipe(
-            prompt=prompt + ", highly detailed, realistic, sepia tone",
+            prompt=prompt,
             image=image,
-            mask_image=processed_mask,
-            num_inference_steps=50,
+            mask_image=mask,
+            height=height,
+            width=width,
             guidance_scale=7.5,
-            max_sequence_length=512,
+            num_inference_steps=50,
+            generator=generator,
         ).images[0]
-        return output, processed_mask, None
+        return output
     except Exception as e:
-        return None, None, f"Error during inpainting: {str(e)}"
+        raise gr.Error(f"Error during inpainting: {str(e)}")
 
 
+# Gradio interface setup
 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",
-        )
-        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],
+    gr.Markdown("# FLUX.1 Fill [dev]")
+    edit_image = gr.ImageEditor(
+        label="Upload and draw mask for inpainting",
+        type="pil",
+        sources=["upload", "webcam"],
+        image_mode="RGB",
+        layers=False,
+        brush=gr.Brush(colors=["#FFFFFF"], color_mode="fixed"),
+        height=600,
     )
+    prompt = gr.Textbox(label="Prompt", placeholder="Enter your prompt")
+    run_button = gr.Button("Run")
+    result = gr.Image(label="Result")
+    run_button.click(infer, inputs=[edit_image, prompt], outputs=result)
 
-if __name__ == "__main__":
-    demo.launch()
+# Launch the demo
+demo.launch()