from diffusers import DiffusionPipeline
import torch



HF_MODEl_ID = {"SD-v1.5": "runwayml/stable-diffusion-v1-5", 
               "SD-v2.1": "runwayml/stable-diffusion-v2-1", 
               "SD-XL-v1.0": "stabilityai/stable-diffusion-xl-base-1.0",
               "IF-I-XL-v1.0": "DeepFloyd/IF-I-XL-v1.0"}


def load_huggingface_model(model_name):
    model_id = HF_MODEl_ID[model_name]
    pipe = DiffusionPipeline.from_pretrained(model_id, torch_dtype=torch.float16)
    pipe = pipe.to("cuda")
    return pipe


for name in ["SD-v1.5", "SD-v2.1", "SD-XL-v1.0", "IF-I-XL-v1.0"]:
    load_huggingface_model(name)