import gradio as gr
import cv2
import requests
import numpy as np
from PIL import Image
from io import BytesIO
from rembg import remove
from diffusers import StableDiffusionControlNetPipeline, ControlNetModel
import torch
from diffusers import UniPCMultistepScheduler


def remove_background(input_image: Image.Image, to_grayscale: bool) -> Image.Image:
    output_image = remove(input_image)
    if to_grayscale:
        output_image = convert_to_grayscale(output_image)
    return output_image

def convert_to_grayscale(image: Image.Image) -> Image.Image:
    return image.convert("L")

def canny_image(image: Image.Image) -> Image.Image:
    np_image = np.array(image)
    low_threshold = 100
    high_threshold = 200
    np_image = cv2.Canny(np_image, low_threshold, high_threshold)
    np_image = np_image[:, :, None]
    np_image = np.concatenate([np_image, np_image, np_image], axis=2)
    return Image.fromarray(np_image)

def process_image(input_image: Image.Image, to_grayscale: bool, prompt: str) -> Image.Image:
    output_image = remove_background(input_image, to_grayscale)
    canny_output = canny_image(output_image)

    controlnet = ControlNetModel.from_pretrained("lllyasviel/sd-controlnet-canny", torch_dtype=torch.float32)
    pipe = StableDiffusionControlNetPipeline.from_pretrained(
        "runwayml/stable-diffusion-v1-5", controlnet=controlnet, torch_dtype=torch.float32
    )
    pipe.enable_model_cpu_offload()
    pipe.scheduler = UniPCMultistepScheduler.from_config(pipe.scheduler.config)
    pipe.enable_xformers_memory_efficient_attention()

    generator = torch.manual_seed(2)
    output = pipe(
        prompt,
        canny_output,
        negative_prompt="monochrome, lowres, bad anatomy, worst quality, low quality",
        generator=generator,
        num_inference_steps=20,
    )

    return output.images[0]

image_input = gr.components.Image(label="Input Image")
grayscale_checkbox = gr.components.Checkbox(label="Convert output to grayscale", default=False)
prompt_input = gr.components.Textbox(lines=1, label="Prompt")
image_output = gr.components.Image(label="Output Image", type="pil")

gr.Interface(
    fn=process_image,
    inputs=[image_input, grayscale_checkbox, prompt_input],
    outputs=image_output,
    title="ControlNet",
    description="Upload an image and a prompt to generate an image with the prompt in the style of the input image.",
).launch()