app.py

import spaces
import torch
from diffusers import FluxInpaintPipeline
import gradio as gr
import re
from PIL import Image,ImageFilter

import os
import numpy as np


def convert_to_fit_size(original_width_and_height, maximum_size = 2048):
    width, height =original_width_and_height
    if width <= maximum_size and height <= maximum_size:
        return width,height
    
    if width > height:
        scaling_factor = maximum_size / width
    else:
        scaling_factor = maximum_size / height

    new_width = int(width * scaling_factor)
    new_height = int(height * scaling_factor)
    return new_width, new_height

def adjust_to_multiple_of_32(width: int, height: int):
    width = width - (width % 32)
    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
device = "cuda" if torch.cuda.is_available() else "cpu"
hf_token = os.environ.get("HF_TOKEN")
pipe = FluxInpaintPipeline.from_pretrained("black-forest-labs/FLUX.1-schnell",token=hf_token, torch_dtype=torch.bfloat16).to(device)






def sanitize_prompt(prompt):
  # Allow only alphanumeric characters, spaces, and basic punctuation
  allowed_chars = re.compile(r"[^a-zA-Z0-9\s.,!?-]")
  sanitized_prompt = allowed_chars.sub("", 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,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")
    if not isinstance(image, dict):
        if image2 == None:
             #print("empty mask")
             return image,None
        else:
            image = dict({'background': image, 'layers': [image2]})

    if image2!=None:
         #print("use image2")
         mask = image2
    else:
         if len(image['layers']) == 0:
              #print("empty mask")
              return image
         #print("use layer")
         mask = image['layers'][0]


    def process_inpaint(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
   

        generators = []
        generator = torch.Generator("cuda").manual_seed(seed)
        generators.append(generator)
        
        fit_width,fit_height = convert_to_fit_size(image.size)
        #print(f"fit {width}x{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)
        mask_image = mask_image.convert("RGB")
        
        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,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)

    return output,mask_image
    

def read_file(path: str) -> str:
    with open(path, 'r', encoding='utf-8') as f:
        content = f.read()

    return content

css="""
#col-left {
    margin: 0 auto;
    max-width: 640px;
}
#col-right {
    margin: 0 auto;
    max-width: 640px;
}
.grid-container {
  display: flex;
  align-items: center;
  justify-content: center;
  gap:10px
}

.image {
  width: 128px; 
  height: 128px; 
  object-fit: cover; 
}

.text {
  font-size: 16px;
}
"""

with gr.Blocks(css=css, elem_id="demo-container") as demo:
    with gr.Column():
        gr.HTML(read_file("demo_header.html"))
        gr.HTML(read_file("demo_tools.html"))
    with gr.Row():
                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):
                        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)
                    with gr.Accordion(label="Advanced Settings", open=False):
                        with gr.Row( equal_height=True):
                            strength = gr.Number(value=0.75, minimum=0, maximum=1.0, step=0.01, label="Inpaint strength")
                            seed = gr.Number(value=0, minimum=0, step=1, label="Inpaint seed")
                        models = ["black-forest-labs/FLUX.1-schnell"]
                        inpaint_model = gr.Dropdown(label="modes", choices=models, value="black-forest-labs/FLUX.1-schnell") 
                        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="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,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"],
                    ["examples/00538245.jpg", "examples/expand_mouth_mask.jpg","a beautiful girl,big-smile",0.75,"examples/expand_mouth_mask_result.jpg"],
                    ["examples/00547245_99.jpg", "examples/00547245_99_mask.jpg","a beautiful girl,eyes closed",0.75,"examples/00547245.jpg"],
                    ["examples/00207245_18.jpg", "examples/00207245_18_mask.jpg","a beautiful girl,mouth opened",0.2,"examples/00207245.jpg"]
                         ]
,
                #fn=example_out,
                inputs=[image,image_mask,prompt,strength,image_out],
                #outputs=[test_out],
                #cache_examples=False,
    )
    gr.HTML(
       gr.HTML(read_file("demo_footer.html"))
    )

    if __name__ == "__main__":
        demo.launch()