from diffusers import StableDiffusionControlNetPipeline, ControlNetModel
from diffusers import UniPCMultistepScheduler
import cv2
import gradio as gr
import numpy as np
import torch
from PIL import Image

# Constants
low_threshold = 100
high_threshold = 200

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

# This command loads the individual model components on GPU on-demand. So, we don't
# need to explicitly call pipe.to("cuda").
pipe.enable_model_cpu_offload()

# Generator seed,
generator = torch.manual_seed(0)

def get_canny_filter(image):
    if not isinstance(image, np.ndarray):
        image = np.array(image) 
        
    image = cv2.Canny(image, low_threshold, high_threshold)
    image = image[:, :, None]
    image = np.concatenate([image, image, image], axis=2)
    canny_image = Image.fromarray(image)
    return canny_image


def generate_images(image, prompt):
    canny_image = get_canny_filter(image)
    output = pipe(
        prompt,
        canny_image,
        generator=generator,
        num_images_per_prompt=3
    )
    all_outputs = []
    all_outputs.append(canny_image)
    for image in output.images:
        all_outputs.append(image)
    return all_outputs


gr.Interface(
    generate_images,
    inputs=[
        gr.Image(type="pil"),
        gr.Textbox(
            label="Enter your prompt",
            max_lines=1,
            placeholder="Sandra Oh, best quality, extremely detailed",
        ),
    ],
    outputs=gr.Gallery().style(grid=[2], height="auto"),
    title="Generate controlled outputs with ControlNet and Stable Diffusion. ",
    description="This Space uses Canny edge maps as the additional conditioning.",
    examples=[["input_image_vermeer.png", "Sandra Oh, best quality, extremely detailed"]],
    allow_flagging=False,
).launch(enable_queue=True)