import torch
import random
import gradio as gr
from diffusers import StableDiffusionControlNetPipeline
from annotator.util import resize_image, HWC3

# Load the pipeline
pipe = StableDiffusionControlNetPipeline.from_pretrained("CompVis/stable-diffusion-v1-4").to("cuda")

def process(input_image, prompt, a_prompt, n_prompt, num_samples, image_resolution, ddim_steps, guess_mode, strength, scale, seed, eta, low_threshold, high_threshold):
    with torch.no_grad():
        img = resize_image(HWC3(input_image), image_resolution)
        
        if seed == -1:
            seed = random.randint(0, 65535)
        generator = torch.manual_seed(seed)

        # Generate images using the pipeline
        images = pipe(prompt=prompt + ', ' + a_prompt, num_inference_steps=ddim_steps, guidance_scale=scale, generator=generator, num_images_per_prompt=num_samples).images

        results = [np.array(image) for image in images]
    return results

block = gr.Blocks().queue()
with block:
    with gr.Row():
        gr.Markdown("## Scene Diffusion with ControlNet")
    with gr.Row():
        with gr.Column():
            input_image = gr.Image(source='upload', type="numpy")
            prompt = gr.Textbox(label="Prompt")
            a_prompt = gr.Textbox(label="Additional Prompt")
            n_prompt = gr.Textbox(label="Negative Prompt")
            num_samples = gr.Slider(label="Images", minimum=1, maximum=12, value=1, step=1)
            image_resolution = gr.Slider(label="Image Resolution", minimum=256, maximum=768, value=512, step=64)
            ddim_steps = gr.Slider(label="Steps", minimum=1, maximum=100, value=20, step=1)
            guess_mode = gr.Checkbox(label='Guess Mode', value=False)
            strength = gr.Slider(label="Strength", minimum=0.0, maximum=1.0, value=0.5, step=0.1)
            scale = gr.Slider(label="Scale", minimum=0.1, maximum=30.0, value=10.0, step=0.1)
            seed = gr.Slider(label="Seed", minimum=0, maximum=10000, value=42, step=1)
            eta = gr.Slider(label="ETA", minimum=0.0, maximum=1.0, value=0.0, step=0.1)
            low_threshold = gr.Slider(label="Canny Low Threshold", minimum=1, maximum=255, value=100, step=1)
            high_threshold = gr.Slider(label="Canny High Threshold", minimum=1, maximum=255, value=200, step=1)
            submit = gr.Button("Generate")
        with gr.Column():
            output_image = gr.Gallery(label='Output', show_label=False, elem_id="gallery").style(grid=2, height='auto')
    submit.click(fn=process, inputs=[input_image, prompt, a_prompt, n_prompt, num_samples, image_resolution, ddim_steps, guess_mode, strength, scale, seed, eta, low_threshold, high_threshold], outputs=output_image)

demo = block
demo.launch(share=True)