app.py

import gradio as gr
import numpy as np

import os
from PIL import Image
import requests
from io import BytesIO
import io
import base64

hf_token = os.environ.get("HF_TOKEN")
auth_headers = {"api_token": hf_token}

def convert_mask_image_to_base64_string(mask_image):
    buffer = io.BytesIO()
    mask_image.save(buffer, format="PNG")  # You can choose the format (e.g., "JPEG", "PNG")
    # Encode the buffer in base64
    image_base64_string = base64.b64encode(buffer.getvalue()).decode('utf-8')
    return f",{image_base64_string}" # for some reason the funciton which downloads image from base64 expects prefix of "," which is redundant in the url

def download_image(url):
    response = requests.get(url)
    return Image.open(BytesIO(response.content)).convert("RGB")

def eraser_api_call(image_base64_file, mask_base64_file, mask_type):

    url = "http://engine.prod.bria-api.com/v1/eraser"
    
    payload = {
    "file": image_base64_file,
    "mask_file": mask_base64_file,
    "mask_type": mask_type,
    }
    response = requests.post(url, json=payload, headers=auth_headers)
    response = response.json()
    res_image = download_image(response["result_url"])
    
    return res_image

ratios_map =  {
    0.5:{"width":704,"height":1408},
    0.57:{"width":768,"height":1344},
    0.68:{"width":832,"height":1216},
    0.72:{"width":832,"height":1152},
    0.78:{"width":896,"height":1152},
    0.82:{"width":896,"height":1088},
    0.88:{"width":960,"height":1088},
    0.94:{"width":960,"height":1024},
    1.00:{"width":1024,"height":1024},
    1.13:{"width":1088,"height":960},
    1.21:{"width":1088,"height":896},
    1.29:{"width":1152,"height":896},
    1.38:{"width":1152,"height":832},
    1.46:{"width":1216,"height":832},
    1.67:{"width":1280,"height":768},
    1.75:{"width":1344,"height":768},
    2.00:{"width":1408,"height":704}
}
ratios = np.array(list(ratios_map.keys()))



def get_masked_image(image, image_mask, width, height):
    image_mask = image_mask # inpaint area is white
    image_mask = image_mask.resize((width, height)) # object to remove is white (1)
    image_mask_pil = image_mask
    image = np.array(image.convert("RGB")).astype(np.float32) / 255.0
    image_mask = np.array(image_mask_pil.convert("L")).astype(np.float32) / 255.0
    assert image.shape[0:1] == image_mask.shape[0:1], "image and image_mask must have the same image size"
    masked_image_to_present = image.copy()
    masked_image_to_present[image_mask > 0.5] = (0.5,0.5,0.5)  # set as masked pixel
    image[image_mask > 0.5] = 0.5  # set as masked pixel - s.t. will be grey 
    image = Image.fromarray((image * 255.0).astype(np.uint8))
    masked_image_to_present = Image.fromarray((masked_image_to_present * 255.0).astype(np.uint8))
    return image, image_mask_pil, masked_image_to_present

    
def get_size(init_image):
    w,h=init_image.size
    curr_ratio = w/h
    ind = np.argmin(np.abs(curr_ratio-ratios))
    ratio = ratios[ind]
    chosen_ratio  = ratios_map[ratio]
    w,h = chosen_ratio['width'], chosen_ratio['height']

    return w,h


def read_content(file_path: str) -> str:
    """read the content of target file
    """
    with open(file_path, 'r', encoding='utf-8') as f:
        content = f.read()

    return content

def predict(dict):

    init_image = Image.fromarray(dict['background'][:, :, :3], 'RGB') #dict['background'].convert("RGB")#.resize((1024, 1024))
    mask = Image.fromarray(dict['layers'][0][:,:,3], 'L') #dict['layers'].convert("RGB")#.resize((1024, 1024))
    
    image_base64_file = convert_mask_image_to_base64_string(init_image)
    mask_base64_file = convert_mask_image_to_base64_string(mask)
    
    mask_type = "manual"
    gen_img = eraser_api_call(image_base64_file, mask_base64_file, mask_type)
    
    return gen_img


css = '''
.gradio-container{max-width: 1100px !important}
#image_upload{min-height:400px}
#image_upload [data-testid="image"], #image_upload [data-testid="image"] > div{min-height: 400px}
#mask_radio .gr-form{background:transparent; border: none}
#word_mask{margin-top: .75em !important}
#word_mask textarea:disabled{opacity: 0.3}
.footer {margin-bottom: 45px;margin-top: 35px;text-align: center;border-bottom: 1px solid #e5e5e5}
.footer>p {font-size: .8rem; display: inline-block; padding: 0 10px;transform: translateY(10px);background: white}
.dark .footer {border-color: #303030}
.dark .footer>p {background: #0b0f19}
.acknowledgments h4{margin: 1.25em 0 .25em 0;font-weight: bold;font-size: 115%}
#image_upload .touch-none{display: flex}
@keyframes spin {
    from {
        transform: rotate(0deg);
    }
    to {
        transform: rotate(360deg);
    }
}
#share-btn-container {padding-left: 0.5rem !important; padding-right: 0.5rem !important; background-color: #000000; justify-content: center; align-items: center; border-radius: 9999px !important; max-width: 13rem; margin-left: auto;}
div#share-btn-container > div {flex-direction: row;background: black;align-items: center}
#share-btn-container:hover {background-color: #060606}
#share-btn {all: initial; color: #ffffff;font-weight: 600; cursor:pointer; font-family: 'IBM Plex Sans', sans-serif; margin-left: 0.5rem !important; padding-top: 0.5rem !important; padding-bottom: 0.5rem !important;right:0;}
#share-btn * {all: unset}
#share-btn-container div:nth-child(-n+2){width: auto !important;min-height: 0px !important;}
#share-btn-container .wrap {display: none !important}
#share-btn-container.hidden {display: none!important}
#prompt input{width: calc(100% - 160px);border-top-right-radius: 0px;border-bottom-right-radius: 0px;}
#run_button {
    width: 100%;
    height: 50px;  /* Set a fixed height for the button */
    display: flex;
    align-items: center;
    justify-content: center;
}
#prompt-container{margin-top:-18px;}
#prompt-container .form{border-top-left-radius: 0;border-top-right-radius: 0}
#image_upload{border-bottom-left-radius: 0px;border-bottom-right-radius: 0px}
'''

image_blocks = gr.Blocks(css=css, elem_id="total-container")
with image_blocks as demo:
    with gr.Column(elem_id="col-container"):
        gr.Markdown("## BRIA Eraser")
        gr.HTML('''
          <p style="margin-bottom: 10px; font-size: 94%">
            This is a demo for BRIA Eraser which enables the ability to remove specific elements or objects. 
            The model was trained on licensed data, and so provide full legal liability coverage for copyright and privacy infringement.
          </p>
        ''')
    with gr.Row():
        with gr.Column():
            image = gr.ImageEditor(sources=["upload"], layers=False, transforms=[], brush=gr.Brush(colors=["#000000"], color_mode="fixed"))
            with gr.Row(elem_id="prompt-container", equal_height=True):
                with gr.Column():  # Wrap the button inside a Column
                    btn = gr.Button("Inpaint!", elem_id="run_button")
        
        with gr.Column():
            image_out = gr.Image(label="Output", elem_id="output-img", height=400)

    # Button click will trigger the inpainting function (no prompt required)
    btn.click(fn=predict, inputs=[image], outputs=[image_out], api_name='run')
    

    gr.HTML(
        """
            <div class="footer">
                <p>Model by <a href="https://huggingface.co/diffusers" style="text-decoration: underline;" target="_blank">Diffusers</a> - Gradio Demo by 🤗 Hugging Face
                </p>
            </div>
        """
    )

image_blocks.queue(max_size=25,api_open=False).launch(show_api=False)