import cv2
import gradio as gr
import numpy as np
import torch

from diffusers import ControlNetModel, DiffusionPipeline, StableDiffusionControlNetPipeline
from PIL import Image

low_threshold = 100
high_threshold = 200

def generate(
    prompt, negative_prompt, num_inference_steps, width, height, guidance_scale, seed, input_image
):
    generator = torch.manual_seed(seed)

    if input_image is None:
        pipeline = DiffusionPipeline.from_pretrained("Lykon/DreamShaper")

        return pipeline(
            prompt=prompt,
            negative_prompt=negative_prompt,
            num_inference_steps=num_inference_steps,
            width=width,
            height=height,
            guidance_scale=guidance_scale,
            generator=generator,
        ).images[0]

    image = cv2.Canny(input_image, low_threshold, high_threshold)
    image = image[:, :, None]
    image = np.concatenate([image, image, image], axis=2)
    canny_image = Image.fromarray(image)

    controlnet = ControlNetModel.from_pretrained("lllyasviel/sd-controlnet-canny")
    pipeline = StableDiffusionControlNetPipeline.from_pretrained("Lykon/DreamShaper", controlnet=controlnet)

    return pipeline(
        prompt=prompt,
        negative_prompt=negative_prompt,
        num_inference_steps=num_inference_steps,
        width=width,
        height=height,
        guidance_scale=guidance_scale,
        generator=generator,
        image=canny_image,
    ).images[0]

iface = gr.Interface(
    fn=generate,
    inputs=[
        gr.Textbox(label="Prompt", value=""),
        gr.Textbox(label="Negative Prompt", value=""),
        gr.Slider(label="Sampling Steps", minimum=1, maximum=150, value=30, step=1),
        gr.Slider(label="Width", minimum=64, maximum=2048, value=512, step=1),
        gr.Slider(label="Height", minimum=64, maximum=2048, value=512, step=1),
        gr.Slider(label="CFG Scale", minimum=1, maximum=30, value=9, step=0.5),
        gr.Slider(
            label="Seed",
            info="Refresh the page to generate a new random seed.",
            minimum=0,
            maximum=4294967294,
            step=1,
            randomize=True,
        ),
        gr.Image(label="Input Image", source='upload', type="numpy")
    ],
    outputs="image",
)

iface.launch()