Spaces:
Running
Running
| import gradio as gr | |
| from PIL import Image | |
| import numpy as np | |
| import os,sys | |
| import torch | |
| import cv2 | |
| import base64 | |
| import io | |
| import multiprocessing | |
| import random | |
| import time | |
| from loguru import logger | |
| from share_btn import community_icon_html, loading_icon_html, share_js | |
| from lama_cleaner.model_manager import ModelManager | |
| from lama_cleaner.schema import Config | |
| try: | |
| torch._C._jit_override_can_fuse_on_cpu(False) | |
| torch._C._jit_override_can_fuse_on_gpu(False) | |
| torch._C._jit_set_texpr_fuser_enabled(False) | |
| torch._C._jit_set_nvfuser_enabled(False) | |
| except: | |
| pass | |
| from lama_cleaner.helper import ( | |
| load_img, | |
| numpy_to_bytes, | |
| resize_max_size, | |
| ) | |
| NUM_THREADS = str(multiprocessing.cpu_count()) | |
| # fix libomp problem on windows https://github.com/Sanster/lama-cleaner/issues/56 | |
| os.environ["KMP_DUPLICATE_LIB_OK"] = "True" | |
| os.environ["OMP_NUM_THREADS"] = NUM_THREADS | |
| os.environ["OPENBLAS_NUM_THREADS"] = NUM_THREADS | |
| os.environ["MKL_NUM_THREADS"] = NUM_THREADS | |
| os.environ["VECLIB_MAXIMUM_THREADS"] = NUM_THREADS | |
| os.environ["NUMEXPR_NUM_THREADS"] = NUM_THREADS | |
| if os.environ.get("CACHE_DIR"): | |
| os.environ["TORCH_HOME"] = os.environ["CACHE_DIR"] | |
| HF_TOKEN_SD = os.environ.get('HF_TOKEN_SD') | |
| device = "cuda" if torch.cuda.is_available() else "cpu" | |
| print(f'device = {device}') | |
| 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 get_image_enhancer(scale = 2, device='cuda:0'): | |
| from basicsr.archs.rrdbnet_arch import RRDBNet | |
| from realesrgan import RealESRGANer | |
| from realesrgan.archs.srvgg_arch import SRVGGNetCompact | |
| from gfpgan import GFPGANer | |
| realesrgan_model = RRDBNet(num_in_ch=3, num_out_ch=3, num_feat=64, | |
| num_block=23, num_grow_ch=32, scale=4 | |
| ) | |
| netscale = scale | |
| model_realesrgan = 'https://github.com/xinntao/Real-ESRGAN/releases/download/v0.1.0/RealESRGAN_x4plus.pth' | |
| upsampler = RealESRGANer( | |
| scale=netscale, | |
| model_path=model_realesrgan, | |
| model=realesrgan_model, | |
| tile=0, | |
| tile_pad=10, | |
| pre_pad=0, | |
| half=False if device=='cpu' else True, | |
| device=device | |
| ) | |
| model_GFPGAN = 'https://github.com/TencentARC/GFPGAN/releases/download/v1.3.0/GFPGANv1.3.pth' | |
| img_enhancer = GFPGANer( | |
| model_path=model_GFPGAN, | |
| upscale=scale, | |
| arch='clean', | |
| channel_multiplier=2, | |
| bg_upsampler=upsampler, | |
| device=device | |
| ) | |
| return img_enhancer | |
| image_enhancer = None | |
| if sys.platform == 'linux' and 0==1: | |
| image_enhancer = get_image_enhancer(scale = 1, device=device) | |
| model = None | |
| def model_process(image, mask, img_enhancer): | |
| global model,image_enhancer | |
| if mask.shape[0] == image.shape[1] and mask.shape[1] == image.shape[0] and mask.shape[0] != mask.shape[1]: | |
| # rotate image | |
| image = np.transpose(image[::-1, ...][:, ::-1], axes=(1, 0, 2))[::-1, ...] | |
| original_shape = image.shape | |
| interpolation = cv2.INTER_CUBIC | |
| size_limit = 1080 | |
| if size_limit == "Original": | |
| size_limit = max(image.shape) | |
| else: | |
| size_limit = int(size_limit) | |
| config = Config( | |
| ldm_steps=25, | |
| ldm_sampler='plms', | |
| zits_wireframe=True, | |
| hd_strategy='Original', | |
| hd_strategy_crop_margin=196, | |
| hd_strategy_crop_trigger_size=1280, | |
| hd_strategy_resize_limit=2048, | |
| prompt='', | |
| use_croper=False, | |
| croper_x=0, | |
| croper_y=0, | |
| croper_height=512, | |
| croper_width=512, | |
| sd_mask_blur=5, | |
| sd_strength=0.75, | |
| sd_steps=50, | |
| sd_guidance_scale=7.5, | |
| sd_sampler='ddim', | |
| sd_seed=42, | |
| cv2_flag='INPAINT_NS', | |
| cv2_radius=5, | |
| ) | |
| if config.sd_seed == -1: | |
| config.sd_seed = random.randint(1, 999999999) | |
| logger.info(f"Origin image shape_0_: {original_shape} / {size_limit}") | |
| image = resize_max_size(image, size_limit=size_limit, interpolation=interpolation) | |
| logger.info(f"Resized image shape_1_: {image.shape}") | |
| logger.info(f"mask image shape_0_: {mask.shape} / {type(mask)}") | |
| mask = resize_max_size(mask, size_limit=size_limit, interpolation=interpolation) | |
| logger.info(f"mask image shape_1_: {mask.shape} / {type(mask)}") | |
| if model is None: | |
| return None | |
| res_np_img = model(image, mask, config) | |
| torch.cuda.empty_cache() | |
| image = Image.open(io.BytesIO(numpy_to_bytes(res_np_img, 'png'))) | |
| if image_enhancer is not None and img_enhancer: | |
| start = time.time() | |
| input_img_rgb = np.array(image) | |
| input_img_bgr = input_img_rgb[...,[2,1,0]] | |
| _, _, enhance_img = image_enhancer.enhance(input_img_bgr, has_aligned=False, | |
| only_center_face=False, paste_back=True) | |
| input_img_rgb = enhance_img[...,[2,1,0]] | |
| img_enhance = Image.fromarray(np.uint8(input_img_rgb)) | |
| image = img_enhance | |
| log_info = f"image_enhancer_: {(time.time() - start) * 1000}ms, {res_np_img.shape} " | |
| logger.info(log_info) | |
| return image | |
| def resize_image(pil_image, new_width=400): | |
| width, height = pil_image.size | |
| new_height = int(height*(new_width/width)) | |
| pil_image = pil_image.resize((new_width, new_height)) | |
| return pil_image | |
| model = ModelManager( | |
| name='lama', | |
| device=device, | |
| ) | |
| image_type = 'pil' # filepath' | |
| def predict(input, img_enhancer): | |
| if input is None: | |
| return None | |
| if image_type == 'filepath': | |
| # input: {'image': '/tmp/tmp8mn9xw93.png', 'mask': '/tmp/tmpn5ars4te.png'} | |
| origin_image_bytes = open(input["image"], 'rb').read() | |
| print(f'origin_image_bytes = ', type(origin_image_bytes), len(origin_image_bytes)) | |
| image, _ = load_img(origin_image_bytes) | |
| mask, _ = load_img(open(input["mask"], 'rb').read(), gray=True) | |
| elif image_type == 'pil': | |
| # input: {'image': pil, 'mask': pil} | |
| image_pil = input['image'] | |
| mask_pil = input['mask'] | |
| image = np.array(image_pil) | |
| mask = np.array(mask_pil.convert("L")) | |
| output = model_process(image, mask, img_enhancer) | |
| return output, [resize_image(image_pil, new_width=400), resize_image(output, new_width=400)], gr.update(visible=True) | |
| css = ''' | |
| .container {max-width: 1150px; margin: auto;padding-top: 1.5rem} | |
| #begin-btn {color: blue; font-size:20px;} | |
| #work-container {min-width: min(160px, 100%) !important;flex-grow: 0 !important} | |
| #op-container{margin: 0 auto; text-align: center;width:fit-content;min-width: min(150px, 100%);flex-grow: 0; flex-wrap: nowrap;} | |
| #erase-btn-container{margin: 0 auto; text-align: center;width:150px;border-width:3px;border-color:#2c9748} | |
| #erase-btn {padding:0;} | |
| #enhancer-checkbox{width:520px} | |
| #enhancer-tip{width:450px} | |
| #enhancer-tip-div{text-align: left} | |
| #image_output{margin: 0 auto; text-align: center;width:640px} | |
| #download-container{margin: 0 auto; text-align: center;width:fit-content; min-width: min(150px, 100%);flex-grow: 0; flex-wrap: nowrap;} | |
| #download-btn-container{margin: 0 auto; text-align: center;width: 100px;border-width:1px;border-color:#2c9748} | |
| #download-btn {padding:0;} | |
| #share-container{margin: 0 auto; text-align: center;width:fit-content; min-width: min(150px, 100%);flex-grow: 0; flex-wrap: nowrap;} | |
| #image_upload .touch-none{display: flex} | |
| @keyframes spin { | |
| from { | |
| transform: rotate(0deg); | |
| } | |
| to { | |
| transform: rotate(360deg); | |
| } | |
| } | |
| #share-btn-container { | |
| display: flex; padding-left: 0.5rem !important; padding-right: 0.5rem !important; background-color: #000000; justify-content: center; align-items: center; border-radius: 9999px !important; width: 13rem; | |
| } | |
| #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.25rem !important; padding-bottom: 0.25rem !important; | |
| } | |
| #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; | |
| } | |
| ''' | |
| start_cleaner = """async() => { | |
| function isMobile() { | |
| try { | |
| document.createEvent("TouchEvent"); return true; | |
| } catch(e) { | |
| return false; | |
| } | |
| } | |
| var gradioEl = document.querySelector('body > gradio-app').shadowRoot; | |
| if (!gradioEl) { | |
| gradioEl = document.querySelector('body > gradio-app'); | |
| } | |
| const group1 = gradioEl.querySelectorAll('#group_1')[0]; | |
| const group2 = gradioEl.querySelectorAll('#group_2')[0]; | |
| const image_upload = gradioEl.querySelectorAll('#image_upload')[0]; | |
| const image_output = gradioEl.querySelectorAll('#image_output')[0]; | |
| const image_output_container = gradioEl.querySelectorAll('#image-output-container')[0]; | |
| const data_image = gradioEl.querySelectorAll('#image_upload [data-testid="image"]')[0]; | |
| const data_image_div = gradioEl.querySelectorAll('#image_upload [data-testid="image"] > div')[0]; | |
| image_output_container.setAttribute('style', 'width: 0px; height:0px; display:none;'); | |
| if (isMobile()) { | |
| const group1_width = group1.offsetWidth; | |
| image_upload.setAttribute('style', 'width:' + (group1_width - 13*2) + 'px; min-height:none;'); | |
| data_image.setAttribute('style', 'width: ' + (group1_width - 14*2) + 'px;min-height:none;'); | |
| data_image_div.setAttribute('style', 'width: ' + (group1_width - 14*2) + 'px;min-height:none;'); | |
| image_output.setAttribute('style', 'width: ' + (group1_width - 13*2) + 'px;min-height:none;'); | |
| const enhancer = gradioEl.querySelectorAll('#enhancer-checkbox')[0]; | |
| enhancer.style.display = "none"; | |
| } else { | |
| max_height = 800; | |
| const container = gradioEl.querySelectorAll('.container')[0]; | |
| container.setAttribute('style', 'max-width: 100%;'); | |
| data_image.setAttribute('style', 'height: ' + max_height + 'px'); | |
| data_image_div.setAttribute('style', 'min-height: ' + max_height + 'px'); | |
| } | |
| if (!(gradioEl.parentNode)) { | |
| const share_btn_container = gradioEl.querySelectorAll('#share-btn-container')[0]; | |
| share_btn_container.setAttribute('style', 'width: 0px; height:0px;'); | |
| const share_btn_share_icon = gradioEl.querySelectorAll('#share-btn-share-icon')[0]; | |
| share_btn_share_icon.setAttribute('style', 'width: 0px; height:0px;'); | |
| } | |
| group1.style.display = "none"; | |
| group2.style.display = "block"; | |
| window['gradioEl'] = gradioEl; | |
| window['doCheckAction'] = 0; | |
| window['checkAction'] = function checkAction() { | |
| try { | |
| if (window['doCheckAction'] == 0) { | |
| var gallery_items = window['gradioEl'].querySelectorAll('#gallery .gallery-item'); | |
| if (gallery_items && gallery_items.length == 2) { | |
| window.clearInterval(window['checkAction_interval']); | |
| window['doCheckAction'] = 1; | |
| gallery_items[gallery_items.length-1].click(); | |
| } | |
| } | |
| } catch(e) { | |
| } | |
| } | |
| window['checkAction_interval'] = window.setInterval("window.checkAction()", 500); | |
| }""" | |
| download_img = """async() => { | |
| Date.prototype.Format = function (fmt) { | |
| var o = { | |
| "M+": this.getMonth() + 1, | |
| "d+": this.getDate(), | |
| "h+": this.getHours(), | |
| "m+": this.getMinutes(), | |
| "s+": this.getSeconds(), | |
| "q+": Math.floor((this.getMonth() + 3) / 3), | |
| "S": this.getMilliseconds() | |
| }; | |
| if (/(y+)/.test(fmt)) | |
| fmt = fmt.replace(RegExp.$1, (this.getFullYear() + "").substr(4 - RegExp.$1.length)); | |
| for (var k in o) | |
| if (new RegExp("(" + k + ")").test(fmt)) fmt = fmt.replace(RegExp.$1, (RegExp.$1.length == 1) ? (o[k]) : (("00" + o[k]).substr(("" + o[k]).length))); | |
| return fmt; | |
| } | |
| var gradioEl = document.querySelector('body > gradio-app').shadowRoot; | |
| if (!gradioEl) { | |
| gradioEl = document.querySelector('body > gradio-app'); | |
| } | |
| const out_image = gradioEl.querySelectorAll('#image_output img')[0]; | |
| if (out_image) { | |
| var x=new XMLHttpRequest(); | |
| x.open("GET", out_image.src, true); | |
| x.responseType = 'blob'; | |
| x.onload = function(e){ | |
| var url = window.URL.createObjectURL(x.response) | |
| var a = document.createElement('a'); | |
| a.href = url; | |
| a.download = (new Date()).Format("yyyyMMdd_hhmmss"); | |
| a.click(); | |
| } | |
| x.send(); | |
| } | |
| }""" | |
| image_blocks = gr.Blocks(css=css, title='Image Cleaner') | |
| with image_blocks as demo: | |
| with gr.Group(elem_id="group_1", visible=True) as group_1: | |
| with gr.Box(): | |
| with gr.Row(elem_id="gallery_row"): | |
| with gr.Column(elem_id="gallery_col"): | |
| gallery = gr.Gallery(value=['./sample_00.jpg','./sample_00_e.jpg'], show_label=False) | |
| gallery.style(grid=[2], height='500px') | |
| with gr.Row(): | |
| with gr.Column(): | |
| begin_button = gr.Button("Let's GO!", elem_id="begin-btn", visible=True) | |
| with gr.Row(): | |
| with gr.Column(): | |
| gr.HTML(""" | |
| <div style='margin: 0 auto; text-align: center;color:red;'> | |
| <p> | |
| Solemnly promise: this application will not collect any user information and image resources. | |
| </p> | |
| </div> | |
| <div style='margin: 0 auto; text-align: center'> | |
| The model comes from <a href='https://github.com/Sanster/lama-cleaner' target=_blank>[<font style='color:blue;'>Lama</font>]</a>. Thanks! ❤️<br> | |
| <a href='https://huggingface.co' target=_blank>[<font style='color:blue;'>huggingface.co</font>]</a> provides code hosting. Thanks! ❤️ | |
| </div> | |
| """ | |
| ) | |
| with gr.Group(elem_id="group_2", visible=False) as group_2: | |
| with gr.Box(elem_id="work-container"): | |
| with gr.Row(elem_id="input-container"): | |
| with gr.Column(): | |
| image_input = gr.Image(source='upload', elem_id="image_upload",tool='sketch', type=f'{image_type}', | |
| label="Upload", show_label=False).style(mobile_collapse=False) | |
| with gr.Row(elem_id="op-container").style(mobile_collapse=False, equal_height=True): | |
| with gr.Column(elem_id="erase-btn-container"): | |
| erase_btn = gr.Button(value = "Erase(⏬)",elem_id="erase-btn").style( | |
| margin=True, | |
| rounded=(True, True, True, True), | |
| full_width=True, | |
| ).style(width=100) | |
| with gr.Column(elem_id="enhancer-checkbox", visible=True if image_enhancer is not None else False): | |
| enhancer_label = 'Enhanced image(processing is very slow, please check only for blurred images)' | |
| img_enhancer = gr.Checkbox(label=enhancer_label).style(width=150) | |
| with gr.Row(elem_id="output-container"): | |
| with gr.Column(elem_id="image-output-container"): | |
| image_out = gr.Image(elem_id="image_output",label="Result", show_label=False, visible=False) | |
| with gr.Column(): | |
| gallery = gr.Gallery( | |
| label="Generated images", show_label=False, elem_id="gallery" | |
| ).style(grid=[2], height="600px") | |
| with gr.Row(elem_id="download-container", visible=False) as download_container: | |
| with gr.Column(elem_id="download-btn-container") as download_btn_container: | |
| download_button = gr.Button(elem_id="download-btn", value="Save(⏩)") | |
| with gr.Column(elem_id="share-container") as share_container: | |
| with gr.Group(elem_id="share-btn-container"): | |
| community_icon = gr.HTML(community_icon_html, elem_id="community-icon", visible=True) | |
| loading_icon = gr.HTML(loading_icon_html, elem_id="loading-icon", visible=True) | |
| share_button = gr.Button("Share to community", elem_id="share-btn", visible=True) | |
| erase_btn.click(fn=predict, inputs=[image_input, img_enhancer], outputs=[image_out, gallery, download_container]) | |
| download_button.click(None, [], [], _js=download_img) | |
| share_button.click(None, [], [], _js=share_js) | |
| begin_button.click(fn=None, inputs=[], outputs=[group_1, group_2], _js=start_cleaner) | |
| image_blocks.launch(server_name='0.0.0.0') | |