Spaces:
Runtime error
Runtime error
| from .diffusion_utils import build_pipeline | |
| NAME_TO_MODEL = { | |
| "stable-diffusion-v1-4": | |
| { | |
| "model" : "CompVis/stable-diffusion-v1-4", | |
| "unet" : "CompVis/stable-diffusion-v1-4", | |
| "tokenizer" : "openai/clip-vit-large-patch14", | |
| "text_encoder" : "openai/clip-vit-large-patch14", | |
| }, | |
| "stable_diffusion_v2_1": | |
| { | |
| "model" : "stabilityai/stable-diffusion-2-1", | |
| "unet" : "stabilityai/stable-diffusion-2-1", | |
| "tokenizer" : "laion/CLIP-ViT-H-14-laion2B-s32B-b79K", | |
| "text_encoder" : "laion/CLIP-ViT-H-14-laion2B-s32B-b79K", | |
| } | |
| } | |
| def get_model(model_name): | |
| model = NAME_TO_MODEL.get(model_name) | |
| if model is None: | |
| raise ValueError(f"Model name {model_name} not found. Available models: {list(NAME_TO_MODEL.keys())}") | |
| vae, tokenizer, text_encoder, unet = build_pipeline(model["model"], model["tokenizer"], model["text_encoder"], model["unet"]) | |
| return vae, tokenizer, text_encoder, unet |