from pathlib import Path

output = Path("./base_model_images")
if not output.exists():
    output.mkdir(exist_ok=True, parents=True)

prompts = {
    0: "astronaut riding a horse",
    1: "a cute corgi",
    2: "A cinematic shot of a baby racoon wearing an intricate italian priest robe",
    3: "portrait photo of a girl, photograph, highly detailed face, depth of field, moody light, golden hair",
    4: "A photo of beautiful mountain with realistic sunset and blue lake, highly detailed, masterpiece",
}

# base_models = ["sd1.5", "sd2", "realistic", "sdxl", "sdxl-turbo"]
# base_models = ["sdxl-turbo"]
# sd1.5 512, sd2 512, realistic 512, sdxl 1024, sdxl-turbo 512
base_models = {
    "sd1.5": {
        "img-size": 512,
        #     "steps": 40,
        #     "cfg": 1.2,
        # "loras": [
            # {
            #     "lora": "/Users/jeqin/work/code/sd/models/lora_1.5/lcm-sdv15.safetensors",
            #     "trigger_words": ""
            # },
            # {
            #     "lora": "/Users/jeqin/work/code/sd/models/lora_1.5/ColoringBook-sd15.safetensors",
            #     "trigger_words": "Coloring Book, ColoringBookAF"
            # },
            # {
            #     "lora-scale": 0.8,
            #     "lora": "/Users/jeqin/work/code/sd/models/lora_1.5/lego-sd15.safetensors",
            #     "trigger_words": "LEGO Creator"
            # }, {
            #     "lora": "/Users/jeqin/work/code/sd/models/lora_1.5/toyglasses-sd15.safetensors",
            #     "trigger_words": "<lora:toyglasses:1>toyglasses"
            # }, {
            #     "lora": "/Users/jeqin/work/code/sd/models/lora_1.5/song_flower-sd15.safetensors",
            #     "trigger_words": "Song Dynasty flower and bird painting"
            # }, {
            #     "lora": "/Users/jeqin/work/code/sd/models/lora_1.5/3dillu-sd15.safetensors",
            #     "trigger_words": ""
            # }, {
            #     "lora": "/Users/jeqin/work/code/sd/models/lora_1.5/Icons-sd15.safetensors",
            #     "trigger_words": " icons, ios icon app",
            #     "lora-scale": 0.8,
            # },{
            #     "lora": "/Users/jeqin/work/code/sd/models/lora_1.5/CuteCartoon-sd15.safetensors",
            #     "trigger_words": "Cartoon,CuteCartoonAF",
            # },
        # ]
    },
    "sd2": {
        # "steps": 40,
        "img-size": 512,
        # "loras": [
        #     {
        #         "lora": "/Users/jeqin/work/code/sd/models/lora_2.1/ColoringBook-sd21.safetensors",
        #         "trigger_words": "ColoringBookAF",
        #     },
            #         {
            #             "lora": "/Users/jeqin/work/code/sd/models/lora_2.1/TShirtDesign-sd21.safetensors",
            #             "trigger_words": "T Shirt Design, TShirtDesignAF",
            #         }, {
            #             "lora": "/Users/jeqin/work/code/sd/models/lora_2.1/StudioGhibli-sd21.safetensors",
            #             "trigger_words": "Studio Ghibli, StdGBRedmAF",
            #         }, {
            #             "lora": "/Users/jeqin/work/code/sd/models/lora_2.1/3D-sd21.safetensors",
            #             "trigger_words": "3D Render Style, 3DRenderAF",
            #         }, {
            #             "lora": "/Users/jeqin/work/code/sd/models/lora_2.1/Stickers-sd21.safetensors",
            #             "trigger_words": " Sticker",
            #         },
        # ]
    },
    "realistic": {
        "img-size": 512,
        # "steps": 20,
        # "loras": [
        #     {
        #         "lora": "/Users/jeqin/work/code/sd/models/lora_1.5/lego-sd15.safetensors",
        #         "trigger_words": "LEGO Creator",
        #         "lora-scale": 0.8,  # Between 0.6-1.0, recommended to use 0.8.
        #     },
            # {
            #     "lora": "/Users/jeqin/work/code/sd/models/lora_1.5/toyglasses-sd15.safetensors",
            #     "trigger_words": "<lora:toyglasses:1>toyglasses"
            # }, {
            #     "lora": "/Users/jeqin/work/code/sd/models/lora_1.5/song_flower-sd15.safetensors",
            #     "trigger_words": "Song Dynasty flower and bird painting"
            # }, {
            #     "lora": "/Users/jeqin/work/code/sd/models/lora_1.5/3dillu-sd15.safetensors",
            #     "trigger_words": ""
            # }, {
            #     "lora": "/Users/jeqin/work/code/sd/models/lora_1.5/Icons-sd15.safetensors",
            #     "trigger_words": " icons, ios icon app",
            #     "lora-scale": 0.8,
            # }, {
            #     "lora": "/Users/jeqin/work/code/sd/models/lora_1.5/CuteCartoon-sd15.safetensors",
            #     "trigger_words": "Cartoon,CuteCartoonAF",
            # },
        # ],
    },
    "sdxl": {
        # "steps": 40,
        "img-size": 1024,
        # "loras": [
        #     {
        #         "lora": "/Users/jeqin/work/code/sd/models/lora_xl/papercut_sdxl.safetensors",
        #         "trigger_words": "papercut style",
        #     },
            # {
            #     "lora": "/Users/jeqin/work/code/sd/models/lora_xl/toy_face-sdxl.safetensors",
            #     "trigger_words": "toy_face",
            #     "lora-scale": 0.9,
            # }, {
            #     "lora": "/Users/jeqin/work/code/sd/models/lora_xl/Products10k-sdxl.safetensors",
            #     "trigger_words": "",
            # }, {
            #     "lora": "/Users/jeqin/work/code/sd/models/lora_xl/ikea_instructions-sdxl.safetensors",
            #     "trigger_words": "",
            # }, {
            #     "lora": "/Users/jeqin/work/code/sd/models/lora_xl/Lego_sdxl.safetensors",
            #     "trigger_words": "LEGO MiniFig",
            # }, {
            #     "lora": "/Users/jeqin/work/code/sd/models/lora_xl/lego_cinematic_sdxl.safetensors",
            #     "trigger_words": "Lego",
            # }, {
            #     "lora": "/Users/jeqin/work/code/sd/models/lora_xl/Emojis-sdxl.safetensors",
            #     "trigger_words": "Emoji",
            # },  {
            #     "lora": "/Users/jeqin/work/code/sd/models/lora_xl/MoviePoster-sdxl.safetensors",
            #     "trigger_words": "Movie Poster, MoviePosterAF",
            # }, {
            #     "lora": "/Users/jeqin/work/code/sd/models/lora_xl/XboxAvatar-sdxl.safetensors",
            #     "trigger_words": "XBOX AVATAR",
            # },
        # ],
    },
    "sdxl-turbo": {
        "img-size": 512,
        # "steps": 4,
        # "loras": [
        #     {
        #         "lora": "/Users/jeqin/work/code/sd/models/lora_xl/papercut_sdxl.safetensors",
        #         "trigger_words": "papercut style"
        #     },
            # {
            #     "lora": "/Users/jeqin/work/code/sd/models/lora_xl/toy_face-sdxl.safetensors",
            #     "trigger_words": "toy_face",
            #     "lora-scale": 0.9,
            # }, {
            #     "lora": "/Users/jeqin/work/code/sd/models/lora_xl/Products10k-sdxl.safetensors",
            #     "trigger_words": "",
            # }, {
            #     "lora": "/Users/jeqin/work/code/sd/models/lora_xl/ikea_instructions-sdxl.safetensors",
            #     "trigger_words": "",
            # }, {
            #     "lora": "/Users/jeqin/work/code/sd/models/lora_xl/Lego_sdxl.safetensors",
            #     "trigger_words": "LEGO MiniFig",
            # }, {
            #     "lora": "/Users/jeqin/work/code/sd/models/lora_xl/lego_cinematic_sdxl.safetensors",
            #     "trigger_words": "Lego",
            # }, {
            #     "lora": "/Users/jeqin/work/code/sd/models/lora_xl/Emojis-sdxl.safetensors",
            #     "trigger_words": "Emoji",
            # }, {
            #     "lora": "/Users/jeqin/work/code/sd/models/lora_xl/MoviePoster-sdxl.safetensors",
            #     "trigger_words": "Movie Poster, MoviePosterAF",
            # }, {
            #     "lora": "/Users/jeqin/work/code/sd/models/lora_xl/XboxAvatar-sdxl.safetensors",
            #     "trigger_words": "XBOX AVATAR",
            # },
        # ]
    }
}