import spaces
import gradio as gr
import re
from PIL import Image
import flux1_inpaint


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=180)
def process_images(image, image2=None,prompt="a girl",negative_prompt=None,inpaint_model="black-forest-labs/FLUX.1-schnell",strength=0.75):
    if negative_prompt == None:
         negative_prompt = ""
    # I'm not sure when this happen
    if not isinstance(image, dict):
        if image2 == None:
             print("empty mask")
             return image
        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]


    output = flux1_inpaint.process_image(image["background"],mask,prompt,negative_prompt,inpaint_model,strength)
        
    return output
    

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;
}
"""
demo_blocks = gr.Blocks(css=css, elem_id="demo-container")
with demo_blocks as demo:
    with gr.Column():
        gr.HTML(read_file("demo_header.html"))
    with gr.Row():
                with gr.Column():
                    image = gr.ImageEditor(height=1000,sources=['upload','clipboard'],transforms=[],image_mode='RGB', layers=False,  elem_id="image_upload", type="pil", label="Upload",brush=gr.Brush(colors=["#999"], color_mode="fixed"))
                    with gr.Row(elem_id="prompt-container",  equal_height=False):
                        with gr.Row():
                            prompt = gr.Textbox(label="Prompt",value="a eyes closed girl,shut eyes",placeholder="Your prompt (what you want in place of what is erased)", elem_id="prompt")
                            
                    btn = gr.Button("Inpaint!", elem_id="run_button")
                    negative_prompt = gr.Textbox(label="Negative Prompt",placeholder="negative prompt",value="worst quality", elem_id="negative_prompt")
                    
                    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.8, minimum=0, maximum=1.0, step=0.01, label="Inpaint strength")
                            blur_radius = gr.Number(value=25, minimum=0.0, maximum=50.0, step=1, label="Blur Radius")
                            edge_expand = gr.Number(value=8, minimum=0.0, maximum=20.0, step=1, label="Edge Expand")
                    with gr.Row(equal_height=True):
                            models = ["black-forest-labs/FLUX.1-schnell"]
                            inpaint_model = gr.Dropdown(label="modes", choices=models, value="black-forest-labs/FLUX.1-schnell") 
                with gr.Column():
                    image_out = gr.Image(sources=[],label="Output", elem_id="output-img")
                    
            

    btn.click(fn=process_images, inputs=[image, image_mask,prompt,negative_prompt,inpaint_model,strength], outputs =image_out, api_name='infer')
    gr.Examples(
               examples=[["examples/catman.jpg", "examples/catman_mask.jpg","He's wearing a dog face."]]
,
                #fn=predict,
                inputs=[image,image_mask,prompt],
                cache_examples=False,
    )
    gr.HTML(
        """
            <div style="text-align: center;">
                <p>Inpaint Code <a href="https://github.com/opencv/opencv/blob/da3debda6d233af90e421e95700c63fc08b83b75/samples/python/inpaint.py" style="text-decoration: underline;" target="_blank">OpenCV inpaint example</a> - Gradio Demo by 🤗 Hugging Face
                </p>
            </div>
        """
    )

demo_blocks.queue(max_size=25).launch(share=False,debug=True)