from PIL import Image
import controlnet_hinter


#model_suffixes = {"canny","normal","depth","openpose","hed","scribble","mlsd","seg"}


def write_converted_files(original_image, prefix):
    controlnet_hinter.hint_canny(original_image).save(prefix + '_canny.png')
    controlnet_hinter.hint_depth(original_image).save(prefix + '_depth.png')
    controlnet_hinter.hint_fake_scribble(original_image).save(
        prefix + '_scribble.png')
    controlnet_hinter.hint_hed(original_image).save(prefix + '_hed.png')
    controlnet_hinter.hint_hough(original_image).save(prefix + '_mlsd.png')
    controlnet_hinter.hint_normal(original_image).save(prefix + '_normal.png')
    controlnet_hinter.hint_openpose(
        original_image).save(prefix + '_openpose.png')
    # controlnet_hinter.hint_scribble(
    #     original_image).save(prefix + '_scribble.png')
    controlnet_hinter.hint_segmentation(
        original_image).save(prefix + '_seg.png')


if __name__ == '__main__':
    image_types = {'bird', 'human', 'room', 'vermeer'}
    for itype in image_types:
        image = Image.open(f"control_images/{itype}_512x512.png")
        write_converted_files(
            image, prefix=f'control_images/converted/control_{itype}')