Update app.py

app.py

@@ -2,25 +2,26 @@ import torch
 import spaces
 import gradio as gr
 from diffusers import FluxInpaintPipeline
+import random
+import numpy as np
 
 MARKDOWN = """
-# FLUX.1 Inpainting 🔥
+# FLUX.1 Inpainting 🎨
 Shoutout to [Black Forest Labs](https://huggingface.co/black-forest-labs) team for 
 creating this amazing model, and a big thanks to [Gothos](https://github.com/Gothos) 
 for taking it to the next level by enabling inpainting with the FLUX.
 """
 
+MAX_SEED = np.iinfo(np.int32).max
 DEVICE = "cuda" if torch.cuda.is_available() else "cpu"
 
 pipe = FluxInpaintPipeline.from_pretrained(
     "black-forest-labs/FLUX.1-schnell", torch_dtype=torch.bfloat16).to(DEVICE)
 
-
 @spaces.GPU()
-def process(input_image_editor, uploaded_mask, input_text, strength, progress=gr.Progress(track_tqdm=True)):
+def process(input_image_editor, uploaded_mask, input_text, strength, seed, randomize_seed, num_inference_steps, progress=gr.Progress(track_tqdm=True)):
     if not input_text:
-        gr.Info("Please enter a text prompt.")
-        return None
+        raise gr.Error("Please enter a text prompt.")
 
     image = input_image_editor['background']
     if uploaded_mask is None:
@@ -29,29 +30,34 @@ def process(input_image_editor, uploaded_mask, input_text, strength, progress=gr
         mask_image = uploaded_mask
 
     if not image:
-        gr.Info("Please upload an image.")
-        return None
+        raise gr.Error("Please upload an image.")
 
     if not mask_image:
-        gr.Info("Please draw or upload a mask on the image.")
-        return None
+        raise gr.Error("Please draw or upload a mask on the image.")
 
     width, height = image.size
 
-    return pipe(
+    if randomize_seed:
+        seed = random.randint(0, MAX_SEED)
+    generator = torch.Generator(device=DEVICE).manual_seed(seed)
+
+    result = pipe(
         prompt=input_text,
         image=image,
         mask_image=mask_image,
         width=width,
         height=height,
-        strength=strength
+        strength=strength,
+        num_inference_steps=num_inference_steps,
+        generator=generator
     ).images[0]
 
+    return result, mask_image, seed
 
-with gr.Blocks() as demo:
+with gr.Blocks(theme=gr.themes.Soft()) as demo:
     gr.Markdown(MARKDOWN)
     with gr.Row():
-        with gr.Column():
+        with gr.Column(scale=1):
             input_image_editor_component = gr.ImageEditor(
                 label='Image',
                 type='pil',
@@ -66,20 +72,38 @@ with gr.Blocks() as demo:
                 placeholder="Enter your prompt",
                 container=False,
             )
-            strength_slider = gr.Slider(
-                minimum=0.0,
-                maximum=1.0,
-                value=0.7,
-                step=0.01,
-                label="Strength"
-            )
-            with gr.Accordion("Upload a mask", open = False):
-                uploaded_mask_component = gr.Image(label = "Already made mask (black pixels will be preserved, white pixels will be redrawn)", sources = ["upload"], type = "pil")
+            with gr.Accordion("Advanced Settings", open=False):
+                strength_slider = gr.Slider(
+                    minimum=0.0,
+                    maximum=1.0,
+                    value=0.7,
+                    step=0.01,
+                    label="Strength"
+                )
+                seed_number = gr.Number(
+                    label="Seed", 
+                    value=42,
+                    precision=0
+                )
+                randomize_seed = gr.Checkbox(label="Randomize seed", value=True)
+                num_inference_steps = gr.Slider(
+                    minimum=1,
+                    maximum=100,
+                    value=30,
+                    step=1,
+                    label="Number of inference steps"
+                )
+            with gr.Accordion("Upload a mask", open=False):
+                uploaded_mask_component = gr.Image(label="Already made mask (black pixels will be preserved, white pixels will be redrawn)", sources=["upload"], type="pil")
             submit_button_component = gr.Button(
-                value='Submit', variant='primary')
-        with gr.Column():
+                value='Generate', variant='primary')
+        with gr.Column(scale=1):
             output_image_component = gr.Image(
                 type='pil', image_mode='RGB', label='Generated image')
+            with gr.Accordion("Debug Info", open=False):
+                output_mask_component = gr.Image(
+                    type='pil', image_mode='RGB', label='Input mask')
+                output_seed = gr.Number(label="Used Seed")
 
     submit_button_component.click(
         fn=process,
@@ -87,10 +111,15 @@ with gr.Blocks() as demo:
             input_image_editor_component,
             uploaded_mask_component,
             input_text_component,
-            strength_slider
+            strength_slider,
+            seed_number,
+            randomize_seed,
+            num_inference_steps
         ],
         outputs=[
-            output_image_component
+            output_image_component,
+            output_mask_component,
+            output_seed
         ]
     )