models.py

# _AVAILABLE_MODELS = [
#     "stabilityai/stable-diffusion-2-1-base",
#     "Lykon/dreamshaper-8",
#     "Lykon/absolute-reality-1.81",
#     # "stabilityai/sdxl-turbo",
#     # "Lykon/dreamshaper-xl-turbo",
#     # "stabilityai/stable-diffusion-xl-base-1.0",
# ]

# _DEFAULT_MODEL = _AVAILABLE_MODELS[0]

# def generate_model_dict():
#     return {
#         "unet_config": "unet/config.json",
#         "unet": "unet/diffusion_pytorch_model.safetensors",
#         "text_encoder_config": "text_encoder/config.json",
#         "text_encoder": "text_encoder/model.safetensors",
#         "vae_config": "vae/config.json",
#         "vae": "vae/diffusion_pytorch_model.safetensors",
#         "diffusion_config": "scheduler/scheduler_config.json",
#         "tokenizer_vocab": "tokenizer/vocab.json",
#         "tokenizer_merges": "tokenizer/merges.txt",
#     }

# _MODELS = {model: generate_model_dict() for model in _AVAILABLE_MODELS}


MODELS = {
    # See https://huggingface.co/stabilityai/sdxl-turbo for the model details and license
    "stabilityai/sdxl-turbo": {
        "unet_config": "unet/config.json",
        "unet": "unet/diffusion_pytorch_model.safetensors",
        "text_encoder_config": "text_encoder/config.json",
        "text_encoder": "text_encoder/model.safetensors",
        "text_encoder_2_config": "text_encoder_2/config.json",
        "text_encoder_2": "text_encoder_2/model.safetensors",
        "vae_config": "vae/config.json",
        "vae": "vae/diffusion_pytorch_model.safetensors",
        "diffusion_config": "scheduler/scheduler_config.json",
        "tokenizer_vocab": "tokenizer/vocab.json",
        "tokenizer_merges": "tokenizer/merges.txt",
        "tokenizer_2_vocab": "tokenizer_2/vocab.json",
        "tokenizer_2_merges": "tokenizer_2/merges.txt",
    },
    "stabilityai/stable-diffusion-xl-base-1.0":{
        "unet_config": "unet/config.json",
        "unet": "unet/diffusion_pytorch_model.safetensors",
        "text_encoder_config": "text_encoder/config.json",
        "text_encoder": "text_encoder/model.safetensors",
        "text_encoder_2_config": "text_encoder_2/config.json",
        "text_encoder_2": "text_encoder_2/model.safetensors",
        "vae_config": "vae/config.json",
        "vae": "vae/diffusion_pytorch_model.safetensors",
        "diffusion_config": "scheduler/scheduler_config.json",
        "tokenizer_vocab": "tokenizer/vocab.json",
        "tokenizer_merges": "tokenizer/merges.txt",
        "tokenizer_2_vocab": "tokenizer_2/vocab.json",
        "tokenizer_2_merges": "tokenizer_2/merges.txt",

    },
    # See https://huggingface.co/stabilityai/stable-diffusion-2-1-base for the model details and license
    "stabilityai/stable-diffusion-2-1-base": {
        "unet_config": "unet/config.json",
        "unet": "unet/diffusion_pytorch_model.safetensors",
        "text_encoder_config": "text_encoder/config.json",
        "text_encoder": "text_encoder/model.safetensors",
        "vae_config": "vae/config.json",
        "vae": "vae/diffusion_pytorch_model.safetensors",
        "diffusion_config": "scheduler/scheduler_config.json",
        "tokenizer_vocab": "tokenizer/vocab.json",
        "tokenizer_merges": "tokenizer/merges.txt",
    },
    "runwayml/stable-diffusion-v1-5": {
        "unet_config": "unet/config.json",
        "unet": "unet/diffusion_pytorch_model.safetensors",
        "text_encoder_config": "text_encoder/config.json",
        "text_encoder": "text_encoder/model.safetensors",
        "vae_config": "vae/config.json",
        "vae": "vae/diffusion_pytorch_model.safetensors",
        "diffusion_config": "scheduler/scheduler_config.json",
        "tokenizer_vocab": "tokenizer/vocab.json",
        "tokenizer_merges": "tokenizer/merges.txt",
    },
    # "Lykon/dreamshaper-8": {
    #     "unet_config": "unet/config.json",
    #     "unet": "unet/diffusion_pytorch_model.safetensors",
    #     "text_encoder_config": "text_encoder/config.json",
    #     "text_encoder": "text_encoder/model.safetensors",
    #     "vae_config": "vae/config.json",
    #     "vae": "vae/diffusion_pytorch_model.safetensors",
    #     "diffusion_config": "scheduler/scheduler_config.json",
    #     "tokenizer_vocab": "tokenizer/vocab.json",
    #     "tokenizer_merges": "tokenizer/merges.txt",
    # },
    "SG161222/Realistic_Vision_V3.0_VAE": {
        "unet_config": "unet/config.json",
        "unet": "unet/diffusion_pytorch_model.safetensors",
        "text_encoder_config": "text_encoder/config.json",
        "text_encoder": "text_encoder/model.safetensors",
        "vae_config": "vae/config.json",
        "vae": "vae/diffusion_pytorch_model.safetensors",
        "diffusion_config": "scheduler/scheduler_config.json",
        "tokenizer_vocab": "tokenizer/vocab.json",
        "tokenizer_merges": "tokenizer/merges.txt",
    }
}
_AVAILABLE_MODELS = list(MODELS.keys())