Spaces:
Running
on
Zero
Running
on
Zero
File size: 614 Bytes
3427608 |
1 2 3 4 5 6 7 8 9 10 11 12 13 14 15 16 17 18 19 20 |
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) |