import os # os.environ["CUDA_VISIBLE_DEVICES"] = "2" import tensorflow as tf import gradio as gr import tensorflow_hub as hub import random import time import PIL.Image from PIL import Image import numpy as np import requests from io import BytesIO # from diffusers import StableDiffusionUpscalePipeline from simple_lama_inpainting import SimpleLama import torch from shutil import copyfile from PowerPaint import app import argparse def pil_to_binary_mask(pil_image, threshold=0): np_image = np.array(pil_image) grayscale_image = Image.fromarray(np_image).convert("L") binary_mask = np.array(grayscale_image) > threshold mask = np.zeros(binary_mask.shape, dtype=np.uint8) for i in range(binary_mask.shape[0]): for j in range(binary_mask.shape[1]): if binary_mask[i,j] == True : mask[i,j] = 1 mask = (mask*255).astype(np.uint8) output_mask = Image.fromarray(mask) return output_mask def tensor_to_image(tensor): tensor = tensor*255 tensor = np.array(tensor, dtype=np.uint8) if np.ndim(tensor)>3: assert tensor.shape[0] == 1 tensor = tensor[0] return PIL.Image.fromarray(tensor) def load_img(path_to_img): max_dim = 512 img = tf.io.read_file(path_to_img) img = tf.image.decode_image(img, channels=3) img = tf.image.convert_image_dtype(img, tf.float32) shape = tf.cast(tf.shape(img)[:-1], tf.float32) long_dim = max(shape) scale = max_dim / long_dim new_shape = tf.cast(shape * scale, tf.int32) img = tf.image.resize(img, new_shape) img = img[tf.newaxis, :] return img # Do main logic (simple version) def start_stylize_simple(img, style_img): # global hub_model hub_model = hub.load('https://tfhub.dev/google/magenta/arbitrary-image-stylization-v1-256/2') # Save to disk, put random number as a ID to avoid collision ID = int(time.time()) img.save(filepath + f'/tmp/tmp_image-{ID}.jpg') style_img.save(filepath + f'/tmp/tmp_style_image-{ID}.jpg') # Load the input images. content_image = load_img(filepath + f'/tmp/tmp_image-{ID}.jpg') style_image = load_img(filepath + f'/tmp/tmp_style_image-{ID}.jpg') stylized_image = hub_model(tf.constant(content_image), tf.constant(style_image))[0] tensor_to_image(stylized_image).save(filepath + f'/tmp/result-{ID}.jpg') return filepath + f'/tmp/result-{ID}.jpg' def background_remove(img): from rembg import new_session from rembg import remove session = new_session('isnet-general-use') # Save to disk, put random number as a ID to avoid collision ID = int(time.time()) img.save(filepath + f'/tmp/tmp_image-{ID}.jpg') with open(filepath + f'/tmp/tmp_image-{ID}.jpg', 'rb') as i: with open(filepath + f'/tmp/tmp_result-{ID}.jpg', 'wb') as o: input = i.read() output = remove(input, session = session) o.write(output) return filepath + f'/tmp/tmp_result-{ID}.jpg' def object_remove(imgs): ts = int(time.time()) os.mkdir(filepath + f'/tmp/tmp_image-{ts}') os.mkdir(filepath + f'/tmp/tmp_mask-{ts}') os.mkdir(filepath + f'/tmp/tmp_output-{ts}') img = imgs["background"].convert("RGB") mask = pil_to_binary_mask(imgs['layers'][-1].convert("RGB")) img.save(filepath + f'/tmp/tmp_image-{ts}/image.png') mask.save(filepath + f'/tmp/tmp_mask-{ts}/image.png') simple_lama = SimpleLama() img_path = filepath + f'/tmp/tmp_image-{ts}/image.png' mask_path = filepath + f'/tmp/tmp_mask-{ts}/image.png' image = Image.open(img_path) mask = Image.open(mask_path).convert('L') result = simple_lama(image, mask) result.save(f"{filepath}/tmp/tmp_output-{ts}/image.png") # os.system(f'simple_lama {filepath}/tmp/tmp_image-{ts}/image.png {filepath}/tmp/tmp_mask-{ts}/image.png {filepath}/tmp/tmp_output-{ts}/image.png') # os.system(f'iopaint run --model=lama --device=cuda --image={filepath}/tmp/tmp_image-{ts} --mask={filepath}/tmp/tmp_mask-{ts} --output={filepath}/tmp/tmp_output-{ts}') # filename = os.listdir(filepath + f'/tmp/tmp_output-{ts}')[0] return filepath + f'/tmp/tmp_output-{ts}/image.png' def upscale(img): #, prompt, upscale_radio): # Save to disk, put random number as a ID to avoid collision ID = int(time.time()) img.save(filepath + f'/tmp/tmp_image-{ID}.jpg') if False: #upscale_radio == 'Stable Diffusion x4 upscaler': # load model and scheduler model_id = "stabilityai/stable-diffusion-x4-upscaler" pipeline = StableDiffusionUpscalePipeline.from_pretrained(model_id, torch_dtype=torch.float16) pipeline = pipeline.to("cuda") # let's download an image #url = "https://huggingface.co/datasets/hf-internal-testing/diffusers-images/resolve/main/sd2-upscale/low_res_cat.png" #response = requests.get(url) low_res_img = Image.open(filepath + f'/tmp/tmp_image-{ID}.jpg') width, height = low_res_img.size low_res_img = low_res_img.resize((128, 128)) # prompt = "a white cat" upscaled_image = pipeline(prompt = prompt, image=low_res_img).images[0] upscaled_image.resize((width, height)).save(filepath + f'/tmp/tmp_result-{ID}.jpg') # Image.open(filepath + f'/tmp/tmp_result-{ID}.jpg').resize((width, height)) else: os.system(f'python3 {filepath}/Real-ESRGAN/inference_realesrgan.py -n RealESRGAN_x4plus -i {filepath}/tmp/tmp_image-{ID}.jpg') copyfile(f'{filepath}/results/tmp_image-{ID}_out.jpg', f'{filepath}/tmp/tmp_result-{ID}.jpg') return filepath + f'/tmp/tmp_result-{ID}.jpg' def in_painting(*args): ID = int(time.time()) global flag global controller if flag == 0: try: controller = app.PowerPaintController(weight_dtype, "./checkpoints/ppt-v1", True, "ppt-v1") flag += 1 except: controller = app.PowerPaintController(weight_dtype, "./checkpoints/ppt-v1", False, "ppt-v1") result = controller.infer(*args)[0][0] result.save(f'{filepath}/tmp/tmp_result-{ID}.jpg') return f'{filepath}/tmp/tmp_result-{ID}.jpg' def radio_click(choice): if choice == "Art style transfer": return [gr.update(visible=True), gr.update(visible=False), gr.update(visible=False), gr.update(visible=False), gr.update(visible=False)] elif choice == "Object erasing": return [gr.update(visible=False), gr.update(visible=True), gr.update(visible=False), gr.update(visible=False), gr.update(visible=False)] elif choice == "In painting": return [gr.update(visible=False), gr.update(visible=False), gr.update(visible=True), gr.update(visible=False), gr.update(visible=False)] elif choice == "Background removal": return [gr.update(visible=False), gr.update(visible=False), gr.update(visible=False), gr.update(visible=True), gr.update(visible=False)] elif choice == "Image upscaling": return [gr.update(visible=False), gr.update(visible=False), gr.update(visible=False), gr.update(visible=False), gr.update(visible=True)] else: return [gr.update(visible=True), gr.update(visible=False), gr.update(visible=False), gr.update(visible=False), gr.update(visible=False)] if __name__ == '__main__': args = argparse.ArgumentParser() args.add_argument("--weight_dtype", type=str, default="float16") args.add_argument("--checkpoint_dir", type=str, default="./checkpoints/ppt-v1") args.add_argument("--version", type=str, default="ppt-v1") args.add_argument("--share", action="store_true") args.add_argument( "--local_files_only", action="store_true", help="enable it to use cached files without requesting from the hub" ) args.add_argument("--port", type=int, default=7860) args = args.parse_args() # initialize the pipeline controller weight_dtype = torch.float16 if args.weight_dtype == "float16" else torch.float32 flag = 0 filepath = os.path.dirname(os.path.abspath(__file__)) physical_devices = tf.config.experimental.list_physical_devices('GPU') for i in physical_devices: tf.config.experimental.set_memory_growth(i, True) os.environ['TFHUB_MODEL_LOAD_FORMAT'] = 'COMPRESSED' # hub_model = hub.load('https://tfhub.dev/google/magenta/arbitrary-image-stylization-v1-256/2') os.environ['GRADIO_TEMP_DIR']="/home/gradio_demos/tmp" with gr.Blocks() as demo: gr.Markdown("# Image2Image Demos") #radio = gr.Radio(["Art style transfer", "Object erasing", "In painting", "Background removal", "Image upscaling"], value="Art style transfer", label = "Choose functionality") radio = gr.Radio(["Art style transfer", "Object erasing", "In painting", "Background removal", "Image upscaling"], value="Art style transfer", label = "Choose functionality") with gr.Column(visible = True) as art_style_transfer_block: gr.Markdown("## Art style transfer") gr.Markdown("### Using [arbitrary-image-stylization-v1](https://www.kaggle.com/models/google/arbitrary-image-stylization-v1/tensorFlow1/256/2) model") with gr.Row(): with gr.Column(): img = gr.Image(sources='upload', type="pil", label='Image to apply art style') img_list = os.listdir(filepath + "/images") img_list_path = [os.path.join(filepath + "/images", image) for image in img_list] example = gr.Examples( inputs=img, examples_per_page=6, examples=img_list_path ) with gr.Column(): style_img = gr.Image(label="Art syle image", sources='upload', type="pil") style_list = os.listdir(filepath + "/style_images") style_list_path = [os.path.join(filepath + "/style_images", style_image) for style_image in style_list] example = gr.Examples( inputs=style_img, examples_per_page=6, examples=style_list_path ) with gr.Column(): # image_out = gr.Image(label="Output", elem_id="output-img", height=400) image_out = gr.Image(label="Stylized image", elem_id="output-img" ,show_share_button=False, type = 'filepath') stylize_button = gr.Button(value="Stylize") with gr.Column(visible = False) as object_erasing_block: gr.Markdown("## Object erasing") gr.Markdown("### Using [lama](https://github.com/enesmsahin/simple-lama-inpainting) model") with gr.Row(): with gr.Column(): imgs4 = gr.ImageEditor(sources='upload', type="pil", label='Image to erase object', interactive=True) img_list = os.listdir(filepath + "/images4") img_list_path = [os.path.join(filepath + "/images4", image) for image in img_list] example = gr.Examples( inputs=imgs4, examples_per_page=6, examples=img_list_path ) with gr.Column(): image_out4 = gr.Image(label="Object removed image" ,show_share_button=False, type = 'filepath') object_remove_button = gr.Button(value="Remove object") with gr.Column(visible = False) as in_painting_block: gr.Markdown("## In painting") gr.Markdown("### Using [Powerpaint](https://github.com/open-mmlab/PowerPaint) model") with gr.Row(): with gr.Column(): #gr.Markdown("### Input image and draw mask") input_image = gr.ImageEditor(sources="upload", type="pil", label='Image to in-paint', interactive=True) img_list = os.listdir(filepath + "/images4") img_list_path = [os.path.join(filepath + "/images4", image) for image in img_list] example = gr.Examples( inputs=input_image, examples_per_page=6, examples=img_list_path ) task = gr.Radio( ["text-guided", "object-removal", "shape-guided", "image-outpainting"], show_label=False, visible=False, ) # Text-guided object inpainting with gr.Tab("Text-guided object inpainting") as tab_text_guided: enable_text_guided = gr.Checkbox( label="Enable text-guided object inpainting", value=True, interactive=False, visible = False ) text_guided_prompt = gr.Textbox(label="Prompt") text_guided_negative_prompt = gr.Textbox(label="negative_prompt") tab_text_guided.select(fn=app.select_tab_text_guided, inputs=None, outputs=task) # currently, we only support controlnet in PowerPaint-v1 if args.version == "ppt-v1": # gr.Markdown("### Controlnet setting") enable_control = gr.Checkbox( label="Enable controlnet", info="Enable this if you want to use controlnet", visible = False ) controlnet_conditioning_scale = gr.Slider( label="controlnet conditioning scale", minimum=0, maximum=1, step=0.05, value=0.5, visible = False ) control_type = gr.Radio(["canny", "pose", "depth", "hed"], label="Control type", visible = False) input_control_image = gr.ImageEditor(sources="upload", type="pil", visible = False) # Object removal inpainting with gr.Tab("Object removal inpainting", visible = False) as tab_object_removal: enable_object_removal = gr.Checkbox( label="Enable object removal inpainting", value=True, info="The recommended configuration for the Guidance Scale is 10 or higher. \ If undesired objects appear in the masked area, \ you can address this by specifically increasing the Guidance Scale.", interactive=False, ) removal_prompt = gr.Textbox(label="Prompt") removal_negative_prompt = gr.Textbox(label="negative_prompt") tab_object_removal.select(fn=app.select_tab_object_removal, inputs=None, outputs=task) # Object image outpainting with gr.Tab("Image outpainting", visible = False) as tab_image_outpainting: enable_object_removal = gr.Checkbox( label="Enable image outpainting", value=True, info="The recommended configuration for the Guidance Scale is 10 or higher. \ If unwanted random objects appear in the extended image region, \ you can enhance the cleanliness of the extension area by increasing the Guidance Scale.", interactive=False, ) outpaint_prompt = gr.Textbox(label="Outpainting_prompt") outpaint_negative_prompt = gr.Textbox(label="Outpainting_negative_prompt") horizontal_expansion_ratio = gr.Slider( label="horizontal expansion ratio", minimum=1, maximum=4, step=0.05, value=1, ) vertical_expansion_ratio = gr.Slider( label="vertical expansion ratio", minimum=1, maximum=4, step=0.05, value=1, ) tab_image_outpainting.select(fn=app.select_tab_image_outpainting, inputs=None, outputs=task) # Shape-guided object inpainting with gr.Tab("Shape-guided object inpainting", visible = False) as tab_shape_guided: enable_shape_guided = gr.Checkbox( label="Enable shape-guided object inpainting", value=True, interactive=False ) shape_guided_prompt = gr.Textbox(label="shape_guided_prompt") shape_guided_negative_prompt = gr.Textbox(label="shape_guided_negative_prompt") fitting_degree = gr.Slider( label="fitting degree", minimum=0, maximum=1, step=0.05, value=1, ) tab_shape_guided.select(fn=app.select_tab_shape_guided, inputs=None, outputs=task) seed = gr.Slider( label="Seed", minimum=0, maximum=2147483647, step=1, randomize=True, ) with gr.Accordion("Advanced options", open=False, visible = False): ddim_steps = gr.Slider(label="Steps", minimum=1, maximum=50, value=45, step=1) scale = gr.Slider( label="Guidance Scale", info="For object removal and image outpainting, it is recommended to set the value at 10 or above.", minimum=0.1, maximum=30.0, value=7.5, step=0.1, ) with gr.Column(): # gr.Markdown("### Inpainting result") # inpaint_result = gr.Gallery(label="Generated image", show_label=True, columns=1) inpaint_result = gr.Image(label="Generated image", elem_id="output-img" ,show_share_button=False, type = 'filepath') #gr.Markdown("### Mask") gallery = gr.Gallery(label="Generated masks", show_label=False, columns=2, visible = False) run_button = gr.Button(value="In-paint") run_button.click( fn=in_painting, #controller.infer, inputs=[ input_image, text_guided_prompt, text_guided_negative_prompt, shape_guided_prompt, shape_guided_negative_prompt, fitting_degree, ddim_steps, scale, seed, task, vertical_expansion_ratio, horizontal_expansion_ratio, outpaint_prompt, outpaint_negative_prompt, removal_prompt, removal_negative_prompt, enable_control, input_control_image, control_type, controlnet_conditioning_scale, ], outputs=[inpaint_result]#, gallery], ) with gr.Column(visible = False) as background_removal_block: gr.Markdown("## Background removal") gr.Markdown("### Using [rembg](https://pypi.org/project/rembg/) model") with gr.Row(): with gr.Column(): img2 = gr.Image(sources='upload', type="pil", label='Image to remove background') img_list = os.listdir(filepath + "/images2") img_list_path = [os.path.join(filepath + "/images2", image) for image in img_list] example = gr.Examples( inputs=img2, examples_per_page=6, examples=img_list_path ) with gr.Column(): # image_out = gr.Image(label="Output", elem_id="output-img", height=400) image_out2 = gr.Image(label="Background removed image", elem_id="output-img" ,show_share_button=False, type = 'filepath') background_remove_button = gr.Button(value="Remove background") with gr.Column(visible = False) as image_upscaling_block: gr.Markdown("## Image upscaling") # gr.Markdown("### Using [Stable Diffusion x4 upscaler](https://huggingface.co/stabilityai/stable-diffusion-x4-upscaler) or [Real-ESRGAN](https://github.com/xinntao/Real-ESRGAN) model") gr.Markdown("### Using [Real-ESRGAN](https://github.com/xinntao/Real-ESRGAN) model") with gr.Row(): with gr.Column(): img3 = gr.Image(sources='upload', type="pil", label='Image to upscale') img_list = os.listdir(filepath + "/images3") img_list_path = [os.path.join(filepath + "/images3", image) for image in img_list] example = gr.Examples( inputs=img3, examples_per_page=6, examples=img_list_path ) # prompt = gr.Textbox(label="Prompt") # upscale_radio = gr.Radio(["Stable Diffusion x4 upscaler", "Real-ESRGAN"], value="Stable Diffusion x4 upscaler", label = "Choose a model") with gr.Column(): # image_out = gr.Image(label="Output", elem_id="output-img", height=400) image_out3 = gr.Image(label="Upscaled image", elem_id="output-img" ,show_share_button=False, type = 'filepath') upscale_button = gr.Button(value="Upscale") stylize_button.click(fn=start_stylize_simple, inputs=[img, style_img], outputs=[image_out], api_name='stylize') background_remove_button.click(fn=background_remove, inputs=[img2], outputs=[image_out2], api_name='background_removal') object_remove_button.click(fn=object_remove, inputs=[imgs4], outputs=[image_out4], api_name='object_removal') upscale_button.click(fn=upscale, inputs=[img3], outputs=[image_out3], api_name='upscale') radio.change(radio_click, radio, [art_style_transfer_block, object_erasing_block, in_painting_block, background_removal_block, image_upscaling_block]) demo.launch(share=False, server_name="0.0.0.0", ssl_verify=False) # demo.launch(share=True)