import numpy as np
import torch
import sys
import os

from diffusers import (
    StableDiffusionControlNetPipeline,
    AutoencoderKL,
    UNet2DConditionModel,
)
from diffusers.utils import load_image

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


def generate_image(seed, control):
    image = pipe(
        prompt=test_prompt,
        negative_prompt=test_negative_prompt,
        width=512,
        height=512,
        generator=torch.Generator(device="cuda").manual_seed(seed),
        image=control,
    ).images[0]
    return image


if __name__ == "__main__":
    output_image_root_folder = "./canny"
    model_id = f"../../control_sd15_canny"
    base_model_id = sys.argv[1] if len(sys.argv) == 2 else None
    canny_edged_image = load_image(
        "https://huggingface.co/takuma104/controlnet_dev/resolve/main/vermeer_canny_edged.png"
    )

    if base_model_id:
        unet = UNet2DConditionModel.from_pretrained(base_model_id, subfolder="unet").to(
            "cuda"
        )
        vae = AutoencoderKL.from_pretrained(base_model_id, subfolder="vae").to("cuda")
        output_types = [
            base_model_id.split("/")[1] + suffix for suffix in ["_unet", "_unet_vae"]
        ]
    else:
        output_types = ["sd15"]

    for output_type in output_types:
        if output_type == "sd15":
            print("SD15 no override config")
            pipe = StableDiffusionControlNetPipeline.from_pretrained(model_id).to(
                "cuda"
            )
        elif output_type.endswith("_unet"):
            print(f"{base_model_id} unet only override config")
            pipe = StableDiffusionControlNetPipeline.from_pretrained(
                model_id, unet=unet
            ).to("cuda")
        elif output_type.endswith("_unet_vae"):
            print(f"{base_model_id} unet & vae override config")
            pipe = StableDiffusionControlNetPipeline.from_pretrained(
                model_id, unet=unet, vae=vae
            ).to("cuda")
        output_folder = f"{output_image_root_folder}/{output_type}"
        os.makedirs(output_folder, exist_ok=True)
        for seed in range(32):
            image = generate_image(seed=seed, control=canny_edged_image)
            image.save(f"{output_folder}/output_{seed:02d}.png")