import io
import requests
import torch
from PIL import Image
import numpy as np

from diffusers import StableDiffusionPipeline, DDIMScheduler, DDIMInverseScheduler


def load_512(image_path, left=0, right=0, top=0, bottom=0, size=512):
    if type(image_path) is str:
        image = np.array(Image.open(image_path))[:, :, :3]
    else:
        image = image_path
    h, w, c = image.shape
    left = min(left, w-1)
    right = min(right, w - left - 1)
    top = min(top, h - left - 1)
    bottom = min(bottom, h - top - 1)
    image = image[top:h-bottom, left:w-right]
    h, w, c = image.shape
    if h < w:
        offset = (w - h) // 2
        image = image[:, offset:offset + h]
    elif w < h:
        offset = (h - w) // 2
        image = image[offset:offset + w]
    image = np.array(Image.fromarray(image).resize((size, size)))
    return image


if __name__=="__main__":
    # configs
    device = "cuda:0"
    num_inference_steps = 50
    guidance_scale = 1.0

    # load model
    model_id = "runwayml/stable-diffusion-v1-5"
    pipe = StableDiffusionPipeline.from_pretrained(model_id, torch_dtype=torch.float32)
    pipe.scheduler = DDIMScheduler.from_config(pipe.scheduler.config)
    pipe.to(device)

    # download image
    # example 1
    r = requests.get("https://github.com/google/prompt-to-prompt/raw/main/example_images/gnochi_mirror.jpeg", timeout=4.0)
    prompt = "a cat sitting next to a mirror"

    # example 2
    # r = requests.get("https://github.com/pix2pixzero/pix2pix-zero/raw/main/assets/test_images/cats/cat_6.png", timeout=4.0)
    # prompt = "a photography of a black and white kitten in a field of daies"
    
    image =  Image.open(io.BytesIO(r.content))
    image = load_512(np.array(image))
    image = Image.fromarray(image)

    # get image latents
    with torch.no_grad():
        init_image = pipe.image_processor.preprocess(image, height=512, width=512)
        init_image = init_image.to(device)
        image_latents = pipe.vae.encode(init_image)['latent_dist'].mean
        image_latents = pipe.vae.config.scaling_factor * image_latents

    # vae reconstruction
    with torch.no_grad():
        image_vae_recon = pipe.vae.decode(image_latents / pipe.vae.config.scaling_factor)['sample']
        image_vae_recon = pipe.image_processor.postprocess(image_vae_recon, output_type='pil')[0]

    # ddim inversion
    pipe.scheduler = DDIMInverseScheduler.from_config(pipe.scheduler.config)
    inversion_latents = pipe(prompt, 
                             num_inference_steps=num_inference_steps, 
                             guidance_scale=guidance_scale, 
                             latents=image_latents, 
                             output_type="latent",
                             )['images']

    # ddim inversion reconstruction
    pipe.scheduler = DDIMScheduler.from_config(pipe.scheduler.config)
    image_ddim_recon = pipe(prompt, 
                            num_inference_steps=num_inference_steps, 
                            guidance_scale=guidance_scale, 
                            latents=inversion_latents,
                            )['images'][0]

    # save the results
    Image.fromarray(np.concatenate([np.array(image), 
                                    np.array(image_vae_recon), 
                                    np.array(image_ddim_recon),
                                    ], axis=1)).save("ddim_inversion.png")