|
#!/usr/bin/env python |
|
|
|
|
|
import os |
|
import sys |
|
import torch |
|
from diffusers import StableDiffusionPipeline |
|
from diffusers import ( |
|
AutoencoderKL, |
|
DDIMScheduler, |
|
DDPMScheduler, |
|
StableDiffusionPipeline, |
|
UNet2DConditionModel, |
|
) |
|
|
|
from transformers import CLIPTextModel, CLIPTokenizer |
|
|
|
|
|
cache_dir = "ssd-cache" |
|
vae_cache_dir = "sd-vae-ft-mse-cache" |
|
os.makedirs(cache_dir, exist_ok=True) |
|
os.makedirs(vae_cache_dir, exist_ok=True) |
|
|
|
pipe = StableDiffusionPipeline.from_pretrained( |
|
"segmind/SSD-1B-fp32", |
|
) |
|
|
|
pipe.save_pretrained(cache_dir) |
|
|
|
|
|
pretrained_vae = AutoencoderKL.from_pretrained( |
|
"madebyollin/sdxl-vae-fp16-fix", subfolder=None, revision=None |
|
) |
|
pretrained_vae.save_pretrained(vae_cache_dir) |
|
|
