# Diffusers' ControlNet Implementation Subjective Evaluation
# https://github.com/takuma104/diffusers/tree/controlnet

import einops
import numpy as np
import torch
import sys
import os

from diffusers import StableDiffusionControlNetPipeline, ControlNetModel, DDIMScheduler

from PIL import Image

test_prompt = "best quality, extremely detailed"
test_negative_prompt = "lowres, bad anatomy, worst quality, low quality"

def generate_image(seed, prompt, negative_prompt, control, guess_mode=False):
    latent = torch.randn((1,4,64,64), device="cpu", generator=torch.Generator(device="cpu").manual_seed(seed)).cuda()
    image = pipe(
        prompt=prompt,
        negative_prompt=negative_prompt,
        guidance_scale=4.0 if guess_mode else 9.0,
        num_inference_steps=50 if guess_mode else 20,
        latents=latent,
        image=control,
        guess_mode=guess_mode,
    ).images[0]
    return image

if __name__ == '__main__':
    model_name = sys.argv[1]
    control_image_folder = '../gen_compare/control_images/converted/'
    output_image_folder = './output_images/diffusers/'
    os.makedirs(output_image_folder, exist_ok=True)

    model_id = f'lllyasviel/sd-controlnet-{model_name}'

    controlnet = ControlNetModel.from_pretrained(model_id)
    pipe = StableDiffusionControlNetPipeline.from_pretrained("runwayml/stable-diffusion-v1-5", 
                                                             revision="non-ema",
                                                             controlnet=controlnet, 
                                                             safety_checker=None).to("cuda")
    pipe.scheduler = DDIMScheduler.from_config(pipe.scheduler.config)

    image_types = {'bird', 'human', 'room', 'vermeer'}

    for image_type in image_types:
        control_image = Image.open(f'{control_image_folder}control_{image_type}_{model_name}.png')
        control = np.array(control_image)[:,:,::-1].copy()
        control = torch.from_numpy(control).float().cuda() / 255.0
        control = torch.stack([control for _ in range(1)], dim=0)
        control = einops.rearrange(control, 'b h w c -> b c h w').clone()

        if model_name == 'normal': # workaround, this should not be necessary
            control = torch.flip(control, dims=[1]) # RGB -> BGR

        for seed in range(4):
            image = generate_image(seed=seed, 
                                   prompt=test_prompt, 
                                   negative_prompt=test_negative_prompt, 
                                   control=control)
            image.save(f'{output_image_folder}output_{image_type}_{model_name}_{seed}.png')
            image = generate_image(seed=seed, 
                                   prompt="", 
                                   negative_prompt="", 
                                   control=control,
                                   guess_mode=True)
            image.save(f'{output_image_folder}output_{image_type}_{model_name}_{seed}_gm.png')
            image = generate_image(seed=seed, 
                                   prompt=test_prompt, 
                                   negative_prompt=test_negative_prompt, 
                                   control=control,
                                   guess_mode=True)
            image.save(f'{output_image_folder}output_{image_type}_{model_name}_{seed}_gm_wp.png')