from diffusers import StableDiffusionXLInpaintPipeline
from PIL import Image, ImageFilter

import gradio as gr
import numpy as np
import time
import math
import random
import imageio
import torch

max_64_bit_int = 2**63 - 1

device = "cuda" if torch.cuda.is_available() else "cpu"
floatType = torch.float16 if torch.cuda.is_available() else torch.float32
variant = "fp16" if torch.cuda.is_available() else None
pipe = StableDiffusionXLInpaintPipeline.from_pretrained("diffusers/stable-diffusion-xl-1.0-inpainting-0.1", torch_dtype = floatType, variant = variant)
pipe = pipe.to(device)

def check(
    source_img,
    prompt,
    uploaded_mask,
    negative_prompt,
    denoising_steps,
    num_inference_steps,
    guidance_scale,
    image_guidance_scale,
    strength,
    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,
    denoising_steps,
    num_inference_steps,
    guidance_scale,
    image_guidance_scale,
    strength,
    randomize_seed,
    seed,
    debug_mode,
    progress = gr.Progress()
):
    check(
        source_img,
        prompt,
        uploaded_mask,
        negative_prompt,
        denoising_steps,
        num_inference_steps,
        guidance_scale,
        image_guidance_scale,
        strength,
        randomize_seed,
        seed,
        debug_mode
    )
    start = time.time()
    progress(0, desc = "Preparing data...")

    if negative_prompt is None:
        negative_prompt = ""

    if denoising_steps is None:
        denoising_steps = 1000

    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 randomize_seed:
        seed = random.randint(0, max_64_bit_int)

    random.seed(seed)
    #pipe = pipe.manual_seed(seed)

    input_image = source_img["image"].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["mask"].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 = 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]

    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 = secondes // 60
    secondes = secondes - (minutes * 60)
    hours = minutes // 60
    minutes = minutes - (hours * 60)
    return [
        output_image,
        "Start again to get a different result. The new image is " + str(output_width) + " pixels large and " + str(output_height) + " pixels high, so an image of " + f'{output_width * output_height:,}' + " pixels. The image have been generated in " + str(hours) + " h, " + str(minutes) + " min, " + str(secondes) + " sec." + limitation,
        input_image,
        mask_image
    ]

def toggle_debug(is_debug_mode):
    if is_debug_mode:
        return [gr.update(visible = True)] * 2
    else:
        return [gr.update(visible = False)] * 2

with gr.Blocks() as interface:
    gr.Markdown(
        """
        <p style="text-align: center;"><b><big><big><big>Inpaint</big></big></big></b></p>
        <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/>
        <br/>
        🚀 Powered by <i>SDXL 1.0</i> artificial intellingence. For illustration purpose, not information purpose. The new content is not based on real information but imagination.
        <br/>
        <ul>
        <li>To change the <b>view angle</b> of your image, I recommend to use <i>Zero123</i>,</li>
        <li>To <b>upscale</b> your image, I recommend to use <i>Ilaria Upscaler</i>,</li>
        <li>To <b>slightly change</b> your image, I recommend to use <i>Image-to-Image SDXL</i>,</li>
        <li>If you need to enlarge the <b>viewpoint</b> of your image, I recommend you to use <i>Uncrop</i>,</li>
        <li>To remove the <b>background</b> of your image, I recommend to use <i>BRIA</i>,</li>
        <li>To make a <b>tile</b> of your image, I recommend to use <i>Make My Image Tile</i>,</li>
        <li>To modify <b>anything else</b> on your image, I recommend to use <i>Instruct Pix2Pix</i>.</li>
        </ul>
        <br/>
        🐌 Slow process... ~1 hour.<br>You can duplicate this space on a free account, it works on CPU and should also run on CUDA.<br/>
        <a href='https://huggingface.co/spaces/Fabrice-TIERCELIN/Inpaint?duplicate=true'><img src='https://img.shields.io/badge/-Duplicate%20Space-blue?labelColor=white&style=flat&logo=data:image/png;base64,iVBORw0KGgoAAAANSUhEUgAAABAAAAAQCAYAAAAf8/9hAAAAAXNSR0IArs4c6QAAAP5JREFUOE+lk7FqAkEURY+ltunEgFXS2sZGIbXfEPdLlnxJyDdYB62sbbUKpLbVNhyYFzbrrA74YJlh9r079973psed0cvUD4A+4HoCjsA85X0Dfn/RBLBgBDxnQPfAEJgBY+A9gALA4tcbamSzS4xq4FOQAJgCDwV2CPKV8tZAJcAjMMkUe1vX+U+SMhfAJEHasQIWmXNN3abzDwHUrgcRGmYcgKe0bxrblHEB4E/pndMazNpSZGcsZdBlYJcEL9Afo75molJyM2FxmPgmgPqlWNLGfwZGG6UiyEvLzHYDmoPkDDiNm9JR9uboiONcBXrpY1qmgs21x1QwyZcpvxt9NS09PlsPAAAAAElFTkSuQmCC&logoWidth=14'></a>
        <br/>
        ⚖️ You can use, modify and share the generated images but not for commercial uses.

        """
    )
    with gr.Column():
        source_img = gr.Image(label = "Your image", source = "upload", tool = "sketch", type = "pil")
        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")
        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)", source = "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")
             denoising_steps = gr.Slider(minimum = 0, maximum = 1000, value = 1000, step = 1, label = "Denoising", info = "lower=irrelevant result, higher=relevant result")
             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.Number(value = 0.99, minimum = 0.01, maximum = 1.0, step = 0.01, label = "Strength", info = "lower=follow the original area, higher=redraw from scratch")
             randomize_seed = gr.Checkbox(label = "\U0001F3B2 Randomize seed (not working, always checked)", 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 (if not randomized)")
             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.Label(label = "Information")
        original_image = gr.Image(label = "Original image", visible = False)
        mask_image = gr.Image(label = "Mask image", visible = False)

    submit.click(toggle_debug, debug_mode, [
        original_image,
        mask_image
    ], queue = False, show_progress = False).then(check, inputs = [
        source_img,
        prompt,
        uploaded_mask,
        negative_prompt,
        denoising_steps,
        num_inference_steps,
        guidance_scale,
        image_guidance_scale,
        strength,
        randomize_seed,
        seed,
        debug_mode
    ], outputs = [], queue = False, show_progress = False).success(inpaint, inputs = [
        source_img,
        prompt,
        uploaded_mask,
        negative_prompt,
        denoising_steps,
        num_inference_steps,
        guidance_scale,
        image_guidance_scale,
        strength,
        randomize_seed,
        seed,
        debug_mode
    ], outputs = [
        inpainted_image,
        information,
        original_image,
        mask_image
    ], scroll_to_output = True)

    gr.Examples(
	    inputs = [
            source_img,
            prompt,
            uploaded_mask,
            negative_prompt,
            denoising_steps,
            num_inference_steps,
            guidance_scale,
            image_guidance_scale,
            strength,
            randomize_seed,
            seed,
            debug_mode
        ],
	    outputs = [
            inpainted_image,
            information,
            original_image,
            mask_image
        ],
        examples = [
                [
                    "./Examples/Example1.png",
                    "A deer, in a forest landscape, ultrarealistic, realistic, photorealistic, 8k",
                    "./Examples/Mask1.webp",
                    "Painting, drawing, cartoon, ugly, malformed, noise, blur, watermark",
                    1000,
                    25,
                    7,
                    1.1,
                    0.99,
                    True,
                    42,
                    False
                ],
                [
                    "./Examples/Example2.webp",
                    "A cat, ultrarealistic, realistic, photorealistic, 8k",
                    "./Examples/Mask2.png",
                    "Painting, drawing, cartoon, ugly, malformed, noise, blur, watermark",
                    1000,
                    25,
                    7,
                    1.1,
                    0.99,
                    True,
                    42,
                    False
                ],
                [
                    "./Examples/Example3.jpg",
                    "An angry old woman, ultrarealistic, realistic, photorealistic, 8k",
                    "./Examples/Mask3.gif",
                    "Painting, drawing, cartoon, ugly, malformed, noise, blur, watermark",
                    1000,
                    25,
                    7,
                    1.5,
                    0.99,
                    True,
                    42,
                    False
                ],
                [
                    "./Examples/Example4.gif",
                    "A laptop, ultrarealistic, realistic, photorealistic, 8k",
                    "./Examples/Mask4.bmp",
                    "Painting, drawing, cartoon, ugly, malformed, noise, blur, watermark",
                    1000,
                    25,
                    7,
                    1.1,
                    0.99,
                    True,
                    42,
                    False
                ],
                [
                    "./Examples/Example5.bmp",
                    "A sand castle, ultrarealistic, realistic, photorealistic, 8k",
                    "./Examples/Mask5.png",
                    "Painting, drawing, cartoon, ugly, malformed, noise, blur, watermark",
                    1000,
                    50,
                    7,
                    1.5,
                    0.5,
                    True,
                    42,
                    False
                ],
            ],
        cache_examples = False,
    )

    interface.queue().launch()