import os
from diffusers import StableDiffusionControlNetPipeline, ControlNetModel, UniPCMultistepScheduler
import torch
import numpy as np
from pipeline_stable_diffusion_controlnet_inpaint import *
from diffusers.pipelines.stable_diffusion.pipeline_stable_diffusion_controlnet import *
import random
#model1
controlnet =  ControlNetModel.from_pretrained("lllyasviel/control_v11p_sd15_openpose", torch_dtype=torch.float16)
pipe = StableDiffusionControlNetInpaintPipeline.from_pretrained(
    "./models/Any-inpainting", controlnet=controlnet, torch_dtype=torch.float16,cache_dir='./models'
)

pipe.scheduler = UniPCMultistepScheduler.from_config(pipe.scheduler.config)
pipe.to('cuda')

# model2
controlnet1 =  ControlNetModel.from_pretrained("lllyasviel/control_v11p_sd15_scribble",torch_dtype=torch.float16)

pipe1 = StableDiffusionControlNetInpaintPipeline.from_pretrained(
    "./models/Any-inpainting", controlnet=controlnet1, torch_dtype=torch.float16
)

pipe1.scheduler = UniPCMultistepScheduler.from_config(pipe.scheduler.config)
pipe1.to('cuda')

def make_inpaint_condition(image, image_mask):
    image = np.array(image.convert("RGB")).astype(np.float32) / 255.0
    image_mask = np.array(image_mask.convert("L"))
    assert image.shape[0:1] == image_mask.shape[0:1], "image and image_mask must have the same image size"
    image[image_mask > 128] = -1.0 # set as masked pixel
    image = np.expand_dims(image, 0).transpose(0, 3, 1, 2)
    image = torch.from_numpy(image)
    return image


def generate_image(prompt:str, negative_prompt:str, openpose_image, original_image, mask_image):

    a = random.randint(10000,90000)
    generator = torch.manual_seed(a)
    # control_image = make_inpaint_condition(original_image, mask_image)
    # images = [openpose_image, control_image]

    image = pipe(
        prompt=prompt,
        # images,
        image=original_image,
        control_image=openpose_image,
        mask_image=mask_image,
        num_inference_steps=20,
        generator=generator,
        negative_prompt=negative_prompt,
        #     controlnet_conditioning_scale=[1.0, 0.8],
    ).images[0]

    return image


def generate_image_sketch(prompt: str, negative_prompt: str, openpose_image, original_image, mask_image):
    b = random.randint(10000, 90000)
    generator = torch.manual_seed(b)
    # control_image = make_inpaint_condition(original_image, mask_image)
    # images = [openpose_image, control_image]

    image = pipe1(
        prompt=prompt,
        # images,
        image=original_image,
        control_image=openpose_image,
        mask_image=mask_image,
        num_inference_steps=20,
        generator=generator,
        negative_prompt=negative_prompt,
        #     controlnet_conditioning_scale=[1.0, 0.8],
    ).images[0]

    return [image]