Spaces:
Sleeping
Sleeping
| import gradio as gr | |
| import numpy as np | |
| import time | |
| import math | |
| import random | |
| import torch | |
| import spaces | |
| from diffusers import StableDiffusionXLInpaintPipeline | |
| from PIL import Image, ImageFilter, ImageEnhance | |
| import PIL.ImageOps | |
| from diffusers.pipelines.stable_diffusion import safety_checker | |
| def sc(self, clip_input, images) : | |
| return images, [False for i in images] | |
| safety_checker.StableDiffusionSafetyChecker.forward = sc | |
| max_64_bit_int = 2**63 - 1 | |
| if torch.cuda.is_available(): | |
| device = "cuda" | |
| floatType = torch.float16 | |
| variant = "fp16" | |
| else: | |
| device = "cpu" | |
| floatType = torch.float32 | |
| variant = None | |
| pipe = StableDiffusionXLInpaintPipeline.from_pretrained("diffusers/stable-diffusion-xl-1.0-inpainting-0.1", torch_dtype = floatType, variant = variant) | |
| pipe = pipe.to(device) | |
| def update_seed(is_randomize_seed, seed): | |
| if is_randomize_seed: | |
| return random.randint(0, max_64_bit_int) | |
| return seed | |
| def toggle_debug(is_debug_mode): | |
| return [gr.update(visible = is_debug_mode)] * 2 | |
| def check( | |
| source_img, | |
| prompt, | |
| uploaded_mask, | |
| negative_prompt, | |
| num_inference_steps, | |
| guidance_scale, | |
| image_guidance_scale, | |
| strength, | |
| denoising_steps, | |
| is_randomize_seed, | |
| seed, | |
| debug_mode, | |
| progress = gr.Progress() | |
| ): | |
| if source_img is None: | |
| raise gr.Error("Please provide an image.") | |
| if prompt is None or prompt == "": | |
| raise gr.Error("Please provide a prompt input.") | |
| def inpaint( | |
| source_img, | |
| prompt, | |
| uploaded_mask, | |
| negative_prompt, | |
| num_inference_steps, | |
| guidance_scale, | |
| image_guidance_scale, | |
| strength, | |
| denoising_steps, | |
| is_randomize_seed, | |
| seed, | |
| debug_mode, | |
| progress = gr.Progress() | |
| ): | |
| check( | |
| source_img, | |
| prompt, | |
| uploaded_mask, | |
| negative_prompt, | |
| num_inference_steps, | |
| guidance_scale, | |
| image_guidance_scale, | |
| strength, | |
| denoising_steps, | |
| is_randomize_seed, | |
| seed, | |
| debug_mode | |
| ) | |
| start = time.time() | |
| progress(0, desc = "Preparing data...") | |
| if negative_prompt is None: | |
| negative_prompt = "" | |
| if num_inference_steps is None: | |
| num_inference_steps = 25 | |
| if guidance_scale is None: | |
| guidance_scale = 7 | |
| if image_guidance_scale is None: | |
| image_guidance_scale = 1.1 | |
| if strength is None: | |
| strength = 0.99 | |
| if denoising_steps is None: | |
| denoising_steps = 1000 | |
| if seed is None: | |
| seed = random.randint(0, max_64_bit_int) | |
| random.seed(seed) | |
| #pipe = pipe.manual_seed(seed) | |
| input_image = source_img["background"].convert("RGB") | |
| original_height, original_width, original_channel = np.array(input_image).shape | |
| output_width = original_width | |
| output_height = original_height | |
| if uploaded_mask is None: | |
| mask_image = source_img["layers"][0].convert("RGB") | |
| else: | |
| mask_image = uploaded_mask.convert("RGB") | |
| mask_image = mask_image.resize((original_width, original_height)) | |
| # Limited to 1 million pixels | |
| if 1024 * 1024 < output_width * output_height: | |
| factor = ((1024 * 1024) / (output_width * output_height))**0.5 | |
| process_width = math.floor(output_width * factor) | |
| process_height = math.floor(output_height * factor) | |
| limitation = " Due to technical limitation, the image have been downscaled and then upscaled."; | |
| else: | |
| process_width = output_width | |
| process_height = output_height | |
| limitation = ""; | |
| # Width and height must be multiple of 8 | |
| if (process_width % 8) != 0 or (process_height % 8) != 0: | |
| if ((process_width - (process_width % 8) + 8) * (process_height - (process_height % 8) + 8)) <= (1024 * 1024): | |
| process_width = process_width - (process_width % 8) + 8 | |
| process_height = process_height - (process_height % 8) + 8 | |
| elif (process_height % 8) <= (process_width % 8) and ((process_width - (process_width % 8) + 8) * process_height) <= (1024 * 1024): | |
| process_width = process_width - (process_width % 8) + 8 | |
| process_height = process_height - (process_height % 8) | |
| elif (process_width % 8) <= (process_height % 8) and (process_width * (process_height - (process_height % 8) + 8)) <= (1024 * 1024): | |
| process_width = process_width - (process_width % 8) | |
| process_height = process_height - (process_height % 8) + 8 | |
| else: | |
| process_width = process_width - (process_width % 8) | |
| process_height = process_height - (process_height % 8) | |
| progress(None, desc = "Processing...") | |
| output_image = inpaint_on_gpu( | |
| seed, | |
| process_width, | |
| process_height, | |
| prompt, | |
| negative_prompt, | |
| input_image, | |
| mask_image, | |
| num_inference_steps, | |
| guidance_scale, | |
| image_guidance_scale, | |
| strength, | |
| denoising_steps | |
| ) | |
| if limitation != "": | |
| output_image = output_image.resize((output_width, output_height)) | |
| if debug_mode == False: | |
| input_image = None | |
| mask_image = None | |
| end = time.time() | |
| secondes = int(end - start) | |
| minutes = math.floor(secondes / 60) | |
| secondes = secondes - (minutes * 60) | |
| hours = math.floor(minutes / 60) | |
| minutes = minutes - (hours * 60) | |
| return [ | |
| output_image, | |
| ("Start again to get a different result. " if is_randomize_seed else "") + "The image has been generated in " + ((str(hours) + " h, ") if hours != 0 else "") + ((str(minutes) + " min, ") if hours != 0 or minutes != 0 else "") + str(secondes) + " sec." + limitation, | |
| input_image, | |
| mask_image | |
| ] | |
| def inpaint_on_gpu2( | |
| seed, | |
| process_width, | |
| process_height, | |
| prompt, | |
| negative_prompt, | |
| input_image, | |
| mask_image, | |
| num_inference_steps, | |
| guidance_scale, | |
| image_guidance_scale, | |
| strength, | |
| denoising_steps | |
| ): | |
| return input_image | |
| def inpaint_on_gpu( | |
| seed, | |
| process_width, | |
| process_height, | |
| prompt, | |
| negative_prompt, | |
| input_image, | |
| mask_image, | |
| num_inference_steps, | |
| guidance_scale, | |
| image_guidance_scale, | |
| strength, | |
| denoising_steps | |
| ): | |
| return pipe( | |
| seeds = [seed], | |
| width = process_width, | |
| height = process_height, | |
| prompt = prompt, | |
| negative_prompt = negative_prompt, | |
| image = input_image, | |
| mask_image = mask_image, | |
| num_inference_steps = num_inference_steps, | |
| guidance_scale = guidance_scale, | |
| image_guidance_scale = image_guidance_scale, | |
| strength = strength, | |
| denoising_steps = denoising_steps, | |
| show_progress_bar = True | |
| ).images[0] | |
| with gr.Blocks() as interface: | |
| gr.HTML( | |
| """ | |
| <h1 style="text-align: center;">Inpaint</h1> | |
| <p style="text-align: center;">Modifies one detail of your image, at any resolution, freely, without account, without watermark, without installation, which can be downloaded</p> | |
| <br/> | |
| """ | |
| ) | |
| with gr.Column(): | |
| source_img = gr.ImageMask(label = "Your image (click on the landscape 🌄 to upload your image; click on the pen 🖌️ to draw the mask)", type = "pil", brush=gr.Brush(colors=["white"], color_mode="fixed")) | |
| prompt = gr.Textbox(label = "Prompt", info = "Describe the subject, the background and the style of image; 77 token limit", placeholder = "Describe what you want to see in the entire image", lines = 2) | |
| with gr.Accordion("Upload a mask", open = False): | |
| uploaded_mask = gr.Image(label = "Already made mask (black pixels will be preserved, white pixels will be redrawn)", sources = ["upload"], type = "pil") | |
| with gr.Accordion("Advanced options", open = False): | |
| negative_prompt = gr.Textbox(label = "Negative prompt", placeholder = "Describe what you do NOT want to see in the entire image", value = "Ugly, malformed, noise, blur, watermark") | |
| num_inference_steps = gr.Slider(minimum = 10, maximum = 100, value = 25, step = 1, label = "Number of inference steps", info = "lower=faster, higher=image quality") | |
| guidance_scale = gr.Slider(minimum = 1, maximum = 13, value = 7, step = 0.1, label = "Classifier-Free Guidance Scale", info = "lower=image quality, higher=follow the prompt") | |
| image_guidance_scale = gr.Slider(minimum = 1, value = 1.1, step = 0.1, label = "Image Guidance Scale", info = "lower=image quality, higher=follow the image") | |
| strength = gr.Slider(value = 0.99, minimum = 0.01, maximum = 1.0, step = 0.01, label = "Strength", info = "lower=follow the original area, higher=redraw from scratch") | |
| denoising_steps = gr.Number(minimum = 0, value = 1000, step = 1, label = "Denoising", info = "lower=irrelevant result, higher=relevant result") | |
| randomize_seed = gr.Checkbox(label = "\U0001F3B2 Randomize seed", value = True, info = "If checked, result is always different") | |
| seed = gr.Slider(minimum = 0, maximum = max_64_bit_int, step = 1, randomize = True, label = "Seed") | |
| debug_mode = gr.Checkbox(label = "Debug mode", value = False, info = "Show intermediate results") | |
| submit = gr.Button("🚀 Inpaint", variant = "primary") | |
| inpainted_image = gr.Image(label = "Inpainted image") | |
| information = gr.HTML() | |
| original_image = gr.Image(label = "Original image", visible = False) | |
| mask_image = gr.Image(label = "Mask image", visible = False) | |
| submit.click(update_seed, inputs = [ | |
| randomize_seed, seed | |
| ], outputs = [ | |
| seed | |
| ], queue = False, show_progress = False).then(toggle_debug, debug_mode, [ | |
| original_image, | |
| mask_image | |
| ], queue = False, show_progress = False).then(check, inputs = [ | |
| source_img, | |
| prompt, | |
| uploaded_mask, | |
| negative_prompt, | |
| num_inference_steps, | |
| guidance_scale, | |
| image_guidance_scale, | |
| strength, | |
| denoising_steps, | |
| randomize_seed, | |
| seed, | |
| debug_mode | |
| ], outputs = [], queue = False, show_progress = False).success(inpaint, inputs = [ | |
| source_img, | |
| prompt, | |
| uploaded_mask, | |
| negative_prompt, | |
| num_inference_steps, | |
| guidance_scale, | |
| image_guidance_scale, | |
| strength, | |
| denoising_steps, | |
| randomize_seed, | |
| seed, | |
| debug_mode | |
| ], outputs = [ | |
| inpainted_image, | |
| information, | |
| original_image, | |
| mask_image | |
| ], scroll_to_output = True) | |
| interface.queue().launch() |