import gradio as gr
import cv2
import numpy as np
import os
import torch
from PIL import Image
from diffusers import StableDiffusionInpaintPipeline


auth_token = os.environ.get("READ_TOKEN") or True


def preview(image, state):
    h, w = image.shape[:2]
    scale_percent = 512 / max([w, h])

    width = int(w * scale_percent)
    height = int(h * scale_percent)
    dim = (width, height)
    resized = cv2.resize(image, dim, interpolation=cv2.INTER_AREA)
    yoff = round((512-height)/2)
    xoff = round((512-width)/2)

    final_image = np.zeros((512, 512, 3), dtype=np.uint8)
    final_image.fill(120)
    final_image[yoff:yoff+height, xoff:xoff+width, :] = resized

    mask_image = np.zeros((512, 512, 3), dtype=np.uint8)
    mask_image.fill(255)
    mask_image[yoff:yoff+height, xoff:xoff+width, :] = 0
    state.clear()
    state.append(mask_image)
    state.append([yoff, xoff, height, width])
    state.append(image)

    return final_image, state


def sd_inpaint(image, prompt, state):
    mask = state[0]
    yoff, xoff, height, width = state[1]
    orig_image = state[2]

    device = "cuda" if torch.cuda.is_available() else "cpu"

    pipe = StableDiffusionInpaintPipeline.from_pretrained(
        "runwayml/stable-diffusion-inpainting",
        revision="fp16",
        torch_dtype=torch.float16,
        use_auth_token=auth_token
    ).to(device)

    output = pipe(prompt=prompt, image=Image.fromarray(image), mask_image=Image.fromarray(mask)).images[0]
    result = np.array(output)
    result[yoff:yoff+height, xoff:xoff+width, :] = orig_image
    result = Image.fromarray(result)

    return result


with gr.Blocks(title='Dreambooth Image Editing and Stable Diffusion Inpainting') as demo:
    state = gr.State([])
    gr.Markdown("# Dreambooth Image Editing and Stable Diffusion Inpainting")
    gr.Markdown("It's difficult to get a good image to use for dreambooth, I do not have many photograhps of myself alone and it's very slow to edit the images (crop the selection, scale it to 512x512 and solve the problem of the background somehow)")
    gr.Markdown("This app uses a combination of image selection, automatic scaling, and  stable diffusion inpainting to speed that process. Follow the next instructions:")
    gr.Markdown("""- Upload an image
- Use the select tool to select the area you want to use for dreambooth
- The image will be resized to 512x512 and fill the rest of with a gray background
- Then click the Inpaint button to use stable diffusion to inpaint the background
- Save the image and use it for dreambooth
    """)
    with gr.Row():
        with gr.Column():
            img_ctr = gr.Image(tool='select')
        with gr.Column():
            output = gr.Image()
    with gr.Row():
        greet_btn = gr.Button("Selection")
    with gr.Row():
        sd_prompt = gr.Textbox(lines=2, label="Stable diffusion prompt")
    with gr.Row():
        final_image = gr.Image()
    with gr.Row():
        stab_btn = gr.Button("Inpaint")

    greet_btn.click(fn=preview, inputs=[img_ctr, state], outputs=[output, state])
    stab_btn.click(fn=sd_inpaint, inputs=[output, sd_prompt, state], outputs=final_image)


demo.launch()