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

@spaces.GPU(duration=420)
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()