Spaces:
Running
Running
| import gradio as gr | |
| import PIL | |
| from PIL import Image | |
| import numpy as np | |
| import os | |
| import uuid | |
| import torch | |
| from torch import autocast | |
| import cv2 | |
| from io import BytesIO | |
| from matplotlib import pyplot as plt | |
| from torchvision import transforms | |
| import io | |
| import logging | |
| import multiprocessing | |
| import random | |
| import time | |
| import imghdr | |
| from pathlib import Path | |
| from typing import Union | |
| from loguru import logger | |
| 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 get_image_ext(img_bytes): | |
| w = imghdr.what("", img_bytes) | |
| if w is None: | |
| w = "jpeg" | |
| return w | |
| def read_content(file_path): | |
| """read the content of target file | |
| """ | |
| with open(file_path, 'rb') as f: | |
| content = f.read() | |
| return content | |
| model = None | |
| def model_process(image, mask): | |
| global model | |
| 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 #1080 # "Original" | |
| 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) | |
| print(f"Origin image shape_0_: {original_shape} / {size_limit}") | |
| image = resize_max_size(image, size_limit=size_limit, interpolation=interpolation) | |
| print(f"Resized image shape_1_: {image.shape}") | |
| print(f"mask image shape_0_: {mask.shape} / {type(mask)}") | |
| mask = resize_max_size(mask, size_limit=size_limit, interpolation=interpolation) | |
| print(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'))) | |
| return image # image | |
| model = ModelManager( | |
| name='lama', | |
| device=device, | |
| ) | |
| image_type = 'pil' # filepath' | |
| def predict(input): | |
| if image_type == 'filepath': | |
| # input: {'image': '/tmp/tmp8mn9xw93.png', 'mask': '/tmp/tmpn5ars4te.png'} | |
| origin_image_bytes = read_content(input["image"]) | |
| print(f'origin_image_bytes = ', type(origin_image_bytes), len(origin_image_bytes)) | |
| image, _ = load_img(origin_image_bytes) | |
| mask, _ = load_img(read_content(input["mask"]), 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) | |
| return output | |
| css = ''' | |
| .container {max-width: 100%;margin: auto;padding-top: 1.5rem} | |
| .output-image, .input-image, .image-preview {height: 600px !important;object-fit: contain} | |
| #image_upload{min-height:610px} | |
| #image_upload [data-testid="image"], #image_upload [data-testid="image"] > div{min-height: 620px} | |
| #image_output{margin: 0 auto; text-align: center;width:640px} | |
| #mask_radio .gr-form{background:transparent; border: none} | |
| #mask_radio .gr-form{background:transparent; border: none; color:#00ff00} | |
| #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 { | |
| 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; | |
| } | |
| ''' | |
| image_blocks = gr.Blocks(css=css) | |
| with image_blocks as demo: | |
| with gr.Group(): | |
| with gr.Box(): | |
| with gr.Row(): | |
| with gr.Column(): | |
| image = gr.Image(source='upload', elem_id="image_upload",tool='sketch', type=f'{image_type}', label="Upload").style(mobile_collapse=False) | |
| with gr.Row(elem_id="prompt-container").style(mobile_collapse=False, equal_height=True): | |
| btn_in = gr.Button("Erase(↓)").style( | |
| margin=True, | |
| rounded=(True, True, True, True), | |
| full_width=True, | |
| ) | |
| with gr.Row(): | |
| with gr.Column(): | |
| image_out = gr.Image(label="Output", elem_id="image_output", visible=True).style(width=640) | |
| btn_in.click(fn=predict, inputs=[image], outputs=[image_out]) | |
| image_blocks.launch() | |