add mask

app.py

@@ -1,4 +1,7 @@
 import spaces
+import torch
+from diffusers import FluxInpaintPipeline
+
 import gradio as gr
 import re
 from PIL import Image
@@ -22,7 +25,7 @@ def process_images(image, image2=None,prompt="a girl",inpaint_model="black-fores
     if not isinstance(image, dict):
         if image2 == None:
              print("empty mask")
-             return image
+             return image,None
         else:
             image = dict({'background': image, 'layers': [image2]})
 
@@ -37,9 +40,25 @@ def process_images(image, image2=None,prompt="a girl",inpaint_model="black-fores
          mask = image['layers'][0]
 
 
+    def process_image(image,mask_image,prompt="a person",model_id="black-forest-labs/FLUX.1-schnell",strength=0.75,seed=0,num_inference_steps=4):
+        if image == None:
+            return None
+        
+        pipe = FluxInpaintPipeline.from_pretrained(model_id, torch_dtype=torch.bfloat16)
+        pipe.to("cuda")
+
+        generators = []
+        generator = torch.Generator("cuda").manual_seed(seed)
+        generators.append(generator)
+        # more parameter see https://huggingface.co/docs/diffusers/api/pipelines/flux#diffusers.FluxInpaintPipeline
+        output = pipe(prompt=prompt, image=image, mask_image=mask_image,generator=generator,strength=strength)
+
+        return output.images[0]
+
+
     output = flux1_inpaint.process_image(image["background"],mask,prompt,inpaint_model,strength,seed)
         
-    return output
+    return output,mask
     
 
 def read_file(path: str) -> str:
@@ -91,12 +110,13 @@ with demo_blocks as demo:
                         id_input=gr.Text(label="Name", visible=False)
                             
                 with gr.Column():
-                    image_out = gr.Image(height=800,sources=[],label="Output", elem_id="output-img",format="jpg")
+                    image_out = gr.Image(height=800,sources=[],label="Output", elem_id="output-img",format="webp")
+                    mask_out = gr.Image(height=800,sources=[],label="Mask", elem_id="mask-img",format="jpeg")
 
                     
             
 
-    btn.click(fn=process_images, inputs=[image, image_mask,prompt,inpaint_model,strength,seed], outputs =image_out, api_name='infer')
+    btn.click(fn=process_images, inputs=[image, image_mask,prompt,inpaint_model,strength,seed], outputs =[image_out,mask_out], api_name='infer')
     gr.Examples(
                examples=[
                     ["images/00547245_99.jpg", "images/00547245_99_mask.jpg","a beautiful girl,eyes closed",0.8,"images/00547245.jpg"],