from controlnet_aux import OpenposeDetector
from diffusers import StableDiffusionControlNetPipeline, ControlNetModel
from diffusers import UniPCMultistepScheduler
import gradio as gr
import torch
from PIL import Image, ImageDraw, ImageFont
import os 
import cv2
import glob
from PIL import Image
import numpy as np
from diffusers.utils import load_image
import random

# Constants
low_threshold = 100
high_threshold = 200

# Models
pose_model = OpenposeDetector.from_pretrained("lllyasviel/ControlNet")
controlnet = ControlNetModel.from_pretrained(
    "lllyasviel/sd-controlnet-openpose"
)
pipe = StableDiffusionControlNetPipeline.from_pretrained(
    "runwayml/stable-diffusion-v1-5", controlnet=controlnet, safety_checker=None
)
pipe = pipe.to("cpu")


def get_pose(image):
    return pose_model(image) 



def generate_an_image_from_text(text, text_size_, width, lenght):
    # Create a blank image
    image = Image.new('RGB', (width, lenght), color = (255, 255, 255))
    # Create a drawing object
    draw = ImageDraw.Draw(image)
    # font def 
    dir_path = ''
    # Get a list of all the font files in the directory
    print("start generation")
    font_files = glob.glob(os.path.join(dir_path, '*.ttf'))
    # Get a list of font paths
    font_paths = []
    for font_file in font_files:
        font_paths.append(font_file)
    # Select a random font
    font_path = random.choice(font_paths)
    #print(font_path)
    font = ImageFont.truetype(font_path, text_size_)
    # Get the text size
    text_size = draw.textsize(text, font)
    # Calculate the x and y positions for the text
    x = (image.width - text_size[0]) / 2
    y = (image.height - text_size[1]) / 2
    # Draw the text on the image
    draw.text((x, y), text, fill=(0, 0, 0), font=font)
    print("end generation")

    return image

def to_Canny(image):
    print("start canny")

    # Let's load the popular vermeer image
    image = np.array(image)
    
    low_threshold = 100
    high_threshold = 200
    
    image = cv2.Canny(image, low_threshold, high_threshold)
    image = image[:, :, None]
    image = np.concatenate([image, image, image], axis=2)
    canny_image = Image.fromarray(image)
    print("end canny")

    return canny_image

def inference(prompt,canny_image,number,seed, steps ):
    print("start inference")

    
    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()
    # xformers
    #pipe.enable_xformers_memory_efficient_attention()
    # Generator seed,
    generator = torch.manual_seed(seed)
    image_ = canny_image
    prompt = prompt
    out_image = pipe(
      prompt, num_inference_steps=steps, generator=generator, image=image_, num_images_per_prompt=number)
    print('end inference')
    return out_image
    


def generation(prompt,text,seed,police_size, lenght, width,number,num_inference_steps):
    img = generate_an_image_from_text(text,police_size,lenght,width)
    img = to_Canny(img)
    output = inference(prompt,img, number,seed,num_inference_steps)
    all_outputs = []
    for image in output.images:
        all_outputs.append(image)
    return all_outputs

gr.Interface(fn=generation, 
             inputs=[gr.Textbox(value="A steampunk Alphabetic Logo, steampunk style, with glowing mecha parts, mecha alphabets, high quality, high res, ultra HD"), gr.Textbox(), gr.Slider(0, 200,value=60), gr.Slider(0, 200, value=90), gr.Slider(0, 1024, value=512), gr.Slider(0, 1024, value=512),
                     gr.Slider(0, 7,value=2, step=1),gr.Slider(0, 20,value=5, step=1)], outputs=gr.Gallery().style(grid=[2], height="auto"), title="Generate a logo using Text ",cache_examples=True, examples=[["A steampunk Alphabetic Logo, steampunk style, with glowing mecha parts, mecha alphabets, high quality, high res, ultra HD", "Logo",60,90,512,512,2,5]]).launch(enable_queue=True)