support donut

app.py

@@ -3,7 +3,7 @@ import torch
 from diffusers import FluxInpaintPipeline
 import gradio as gr
 import re
-from PIL import Image
+from PIL import Image,ImageFilter
 
 import os
 import numpy as np
@@ -28,6 +28,16 @@ def adjust_to_multiple_of_32(width: int, height: int):
     height = height - (height % 32)
     return width, height
 
+def mask_to_donut(mask,size):
+    if size%2 ==0:
+        size+=1
+    dilation_mask = mask.filter(ImageFilter.MaxFilter(size))
+
+    white_img = Image.new('RGB', mask.size, (255,255,255))
+    black_img = Image.new('RGB', mask.size, (0,0,0))
+    white_img.paste(black_img,(0,0),dilation_mask.convert("L"))
+    white_img.paste(mask,(0,0),mask.convert("L"))
+    return white_img
 
 
 dtype = torch.bfloat16
@@ -47,7 +57,7 @@ def sanitize_prompt(prompt):
   return sanitized_prompt
 
 @spaces.GPU(duration=120)
-def process_images(image, image2=None,prompt="a girl",inpaint_model="black-forest-labs/FLUX.1-schnell",strength=0.75,seed=0,progress=gr.Progress(track_tqdm=True)):
+def process_images(image, image2=None,prompt="a girl",inpaint_model="black-forest-labs/FLUX.1-schnell",strength=0.75,seed=0,donut_mask=True,donut_size=32,progress=gr.Progress(track_tqdm=True)):
     # I'm not sure when this happen
     progress(0, desc="start-process-images")
     #print("start-process-images")
@@ -78,9 +88,9 @@ def process_images(image, image2=None,prompt="a girl",inpaint_model="black-fores
         generator = torch.Generator("cuda").manual_seed(seed)
         generators.append(generator)
         
-        width,height = convert_to_fit_size(image.size)
+        fit_width,fit_height = convert_to_fit_size(image.size)
         #print(f"fit {width}x{height}")
-        width,height = adjust_to_multiple_of_32(width,height)
+        width,height = adjust_to_multiple_of_32(fit_width,fit_height)
         #print(f"multiple {width}x{height}")
         image = image.resize((width, height), Image.LANCZOS)
         mask_image = mask_image.resize((width, height), Image.NEAREST)
@@ -89,11 +99,24 @@ def process_images(image, image2=None,prompt="a girl",inpaint_model="black-fores
         output = pipe(prompt=prompt, image=image, mask_image=mask_image,generator=generator,strength=strength,width=width,height=height,
                       guidance_scale=0,num_inference_steps=num_inference_steps,max_sequence_length=256)
 
-        return output.images[0],mask_image
+        return output.images[0],mask_image,image,fit_width,fit_height
+
+    if donut_mask:
+        original_mask = mask
+        mask = mask_to_donut(mask,donut_size)
+    
+    #output,mask_image,image_resized,fit_width,fit_height=image["background"],mask,image["background"],512,512
+    output,mask_image,image_resized,fit_width,fit_height = process_inpaint(image["background"],mask,prompt,inpaint_model,strength,seed)
+    
+    if donut_mask:
+        mask = original_mask.resize(mask_image.size)
+        image_resized.paste(output,(0,0),mask.convert("L"))
+        output = image_resized.resize((fit_width,fit_height),Image.LANCZOS)
+        mask_image = mask.resize(output.size)
+    else:
+        output = output.resize((fit_width,fit_height),Image.LANCZOS)
+        mask_image = mask_image.resize(output.size)
 
-   
-    output,mask_image = process_inpaint(image["background"],mask,prompt,inpaint_model,strength,seed)
-        
     return output,mask_image
     
 
@@ -138,9 +161,10 @@ with gr.Blocks(css=css, elem_id="demo-container") as demo:
                 with gr.Column():
                     image = gr.ImageEditor(height=800,sources=['upload','clipboard'],transforms=[],image_mode='RGB', layers=False,  elem_id="image_upload", type="pil", label="Upload",brush=gr.Brush(colors=["#fff"], color_mode="fixed"))
                     with gr.Row(elem_id="prompt-container",  equal_height=False):
-                        with gr.Row():
-                            prompt = gr.Textbox(label="Prompt",value="a person",placeholder="Your prompt (what you want in place of what is erased)", elem_id="prompt")
-                            
+                        prompt = gr.Textbox(label="Prompt",value="a person",placeholder="Your prompt (what you want in place of what is erased)", elem_id="prompt")
+                    with gr.Row(equal_height=True):
+                        donut_mask = gr.Checkbox(label="Donut Mask",value=False,info="Usually improve result,but slow.Do second example things")
+                        donut_size = gr.Slider(label="Donut Size",minimum=1,maximum=64,step=1,value=32,info="Larger value make  extreamly slow")
                     btn = gr.Button("Inpaint", elem_id="run_button",variant="primary")
                     
                     image_mask = gr.Image(sources=['upload','clipboard'],  elem_id="mask_upload", type="pil", label="Mask_Upload",height=400, value=None)
@@ -159,7 +183,7 @@ with gr.Blocks(css=css, elem_id="demo-container") as demo:
                     
             
 
-    btn.click(fn=process_images, inputs=[image, image_mask,prompt,inpaint_model,strength,seed], outputs =[image_out,mask_out], api_name='infer')
+    btn.click(fn=process_images, inputs=[image, image_mask,prompt,inpaint_model,strength,seed,donut_mask,donut_size], outputs =[image_out,mask_out], api_name='infer')
     gr.Examples(
                examples=[
                     ["examples/00538245.jpg", "examples/normal_mouth_mask.jpg","a beautiful girl,big-smile",0.75,"examples/normal_mouth_mask_result.jpg"],