import sys
import yaml
import os

cnet_models = [
    "p_sd15_canny",
    "p_sd15_depth", 
    "p_sd15_mlsd" ,
    "p_sd15_normalbae", 
    "p_sd15_openpose", 
    "p_sd15_scribble", 
    "p_sd15_seg", 
    "p_sd15_softedge",
    "p_sd15_lineart",
    "p_sd15s2_lineart_anime",
    "e_sd15_shuffle",
    # "p_sd15_inpaint",
    # "e_sd15_ip2p",
    # "u_sd15_tile",
]

processors = ["ref", "diffusers"]

print(f"# ControlNet Generation Comparison")
print("")

def control_image_filenames(model_name):
    with open('./control_images.yaml', 'r') as f:
        d = yaml.safe_load(f)
    return d[model_name]

def print_header():
    print("|seed|conditional image|", end="")
    for proc in processors:
        print(f"{proc}|", end="")
    print("")

    print("|", end="")
    for _ in range(len(processors) + 2):
        print(":---:|", end="")
    print("")

def get_url(path):
    # url = "./"
    url = 'https://huggingface.co/takuma104/controlnet_dev/resolve/main/gen_compare_v11/'
    return url + path

control_image_folder = 'control_images/converted'

for cnet in cnet_models:
    print(f"## control_v11{cnet}")
    print("")
    print_header()

    for i, control_image_fn in enumerate(control_image_filenames(cnet)):
        for seed in range(4):
            print(f"|{seed}|", end="")

            url = get_url(path=f"{control_image_folder}/{control_image_fn}")
            print(f'<a href="{url}"><img width="128" src="{url}"/></a>|', end="")

            for proc in processors:
                url = get_url(path=f"output_images/{proc}/output_{cnet}_{i}_{seed:01d}.png")
                print(f'<a href="{url}"><img width="128" src="{url}"/></a>|', end="")

            print("")


original_image_folder = "control_images"
original_image_filenames = [
    "pexels-sound-on-3760767_512x512.png",
    "vermeer_512x512.png",
    "bird_512x512.png",
]

cnet = "p_sd15_inpaint"
print(f"## control_v11{cnet}")
print("")
print_header()

for i, original_image_fn in enumerate(original_image_filenames):
    for seed in range(4):
        print(f"|{seed}|", end="")

        url = get_url(path=f"{original_image_folder}/{original_image_fn}")
        print(f'<a href="{url}"><img width="128" src="{url}"/></a>', end="")
        url = get_url(path=f"{original_image_folder}/mask_512x512.png")
        print(f'<a href="{url}"><img width="128" src="{url}"/></a>|', end="")

        for proc in processors:
            url = get_url(path=f"output_images/{proc}/output_{cnet}_{i}_{seed:01d}.png")
            print(f'<a href="{url}"><img width="128" src="{url}"/></a>|', end="")

        print("")