# !pip install transformers accelerate
from diffusers import StableDiffusionControlNetInpaintPipeline, ControlNetModel, DDIMScheduler
from diffusers.utils import load_image
import numpy as np
import torch
from dw_pose.main import dwpose
from PIL import Image

init_image = load_image(
    "https://huggingface.co/datasets/diffusers/test-arrays/resolve/main/stable_diffusion_inpaint/boy.png"
)
# init_image = init_image.resize((512, 512))
init_image = Image.open("./girl.png")

generator = torch.Generator(device="cpu").manual_seed(1)

mask_image = load_image(
    "https://huggingface.co/datasets/diffusers/test-arrays/resolve/main/stable_diffusion_inpaint/boy_mask.png"
)
# mask_image = mask_image.resize((512, 512))
mask_image = Image.open("./girlmask.png")

print("doing dwpose")

control_image = dwpose(init_image, 512)

print("done dwpose")

control_image.save("thecontrolfucking.png")

print("loading openpose")

controlnet = ControlNetModel.from_pretrained(
    "lllyasviel/control_v11p_sd15_openpose", torch_dtype=torch.float16
)
pipe = StableDiffusionControlNetInpaintPipeline.from_pretrained(
    "runwayml/stable-diffusion-v1-5", controlnet=controlnet, torch_dtype=torch.float16
)

print("loaded openpsoe")

# speed up diffusion process with faster scheduler and memory optimization
pipe.scheduler = DDIMScheduler.from_config(pipe.scheduler.config)

pipe.enable_model_cpu_offload()

pipe.safety_checker = None

init_image_width, init_image_height = init_image.size

from main import resize

new_size = resize(
        init_image_width, init_image_height, 1500)

init_image_resized = init_image.resize((int(new_size['width']), int(new_size['height'])))
mask_image_resized = mask_image.resize((int(new_size['width']), int(new_size['height'])))
control_image_resized = control_image.resize((int(new_size['width']), int(new_size['height'])))

print("ccreating image")

# generate image
image = pipe(
    "a fucking sexy girl naked",
    num_inference_steps=20,
    generator=generator,
    guidance_scale=6.0,
    eta=1.0,
    image=init_image_resized,
    mask_image=mask_image_resized,
    control_image=control_image_resized,
).images[0].save("aha.png")

print("created image and saved")