import spaces 
import gradio as gr
import torch
import modin.pandas as pd
import numpy as np
from diffusers import DiffusionPipeline, DPMSolverSinglestepScheduler, ControlNetModel, StableDiffusionXLControlNetPipeline, AutoencoderKL
import requests 
import cv2 
from uuid import uuid4
import numpy as np 
from PIL import Image

controlnet = ControlNetModel.from_pretrained("diffusers/controlnet-canny-sdxl-1.0", torch_dtype=torch.float16).to("cuda")

pipe = StableDiffusionXLControlNetPipeline.from_pretrained("mann-e/Mann-E_Dreams", controlnet=controlnet, torch_dtype=torch.float16).to("cuda")
pipe.scheduler = DPMSolverSinglestepScheduler.from_config(pipe.scheduler.config, use_karras_sigmas=True)

#pipe.enable_xformers_memory_efficient_attention()
pipe.enable_vae_slicing()

torch.cuda.empty_cache()

@spaces.GPU
def genie (input_image, prompt, negative_prompt, width, height, steps, seed, conditioning_scale):
    #processing the input image
    res = requests.get(input_image)
    image_name = f'tmp_{uuid4()}.png'

    if res.ok:
        with open(image_name, 'wb') as f:
            f.write(res.content)

    # Canny Edge Detection 
    image = cv2.imread(image_name)
    image = np.array(image)
    image = cv2.Canny(image, 100, 200)
    image = image[:, :, None]
    image = np.concatenate([image, image, image], axis=2)
    #cv2.imwrite('canny.png', image)
    image = Image.fromarray(image)
    
    #generating a new image
    generator = np.random.seed(0) if seed == 0 else torch.manual_seed(seed)
    int_image = pipe(prompt=prompt, negative_prompt=negative_prompt, width=width, height=height, generator=generator, num_inference_steps=steps, guidance_scale=5.0, image=image, controlnet_conditioning_scale=conditioning_scale).images[0]
    return int_image
    
gr.Interface(fn=genie, inputs=[gr.Textbox(label='Base Image URL'),
    gr.Textbox(label='What you want the AI to generate. 75 Token Limit.'),
    gr.Textbox(label='What you DO NOT want the AI to generate. 75 Token Limit.'),
    gr.Slider(576, maximum=1280, value=768, step=16, label='Width (can go up to 1280, but for square images maximum is 1024x1024)'), 
    gr.Slider(576, maximum=1280, value=768, step=16, label='Height (can go up to 1280, but for square images maximum is 1024x1024)'), 
    gr.Slider(1, maximum=8, value=6, step=1, label='Number of Iterations'), 
    gr.Slider(minimum=0, step=1, maximum=999999999999999999, randomize=True, label="Seed"),
    gr.Slider(minimum=0, step=0.05, maximum=1, label="Conditioning Scale"),
    ],
    outputs='image', 
    title="Mann-E Dreams w/ ControlNet", 
    description="Mann-E Dreams <br><br><b>WARNING: This model is capable of producing NSFW (Softcore) images.</b><br><br>In case you don't want a base image, just paste link to an image and put conditioning scale to 0", 
    article = "").launch(debug=True, max_threads=80, show_error=True)