Download and save VAE to local directory.
(vae_type: str, output_path: str)
| 170 | |
| 171 | |
| 172 | def download_and_save_vae(vae_type: str, output_path: str): |
| 173 | """Download and save VAE to local directory.""" |
| 174 | from diffusers import AutoencoderDC, AutoencoderKL |
| 175 | |
| 176 | vae_path = os.path.join(output_path, "vae") |
| 177 | os.makedirs(vae_path, exist_ok=True) |
| 178 | |
| 179 | if vae_type == "flux": |
| 180 | print("Downloading FLUX VAE from black-forest-labs/FLUX.1-dev...") |
| 181 | vae = AutoencoderKL.from_pretrained("black-forest-labs/FLUX.1-dev", subfolder="vae") |
| 182 | else: # dc-ae |
| 183 | print("Downloading DC-AE VAE from mit-han-lab/dc-ae-f32c32-sana-1.1-diffusers...") |
| 184 | vae = AutoencoderDC.from_pretrained("mit-han-lab/dc-ae-f32c32-sana-1.1-diffusers") |
| 185 | |
| 186 | vae.save_pretrained(vae_path) |
| 187 | print(f"✓ Saved VAE to {vae_path}") |
| 188 | |
| 189 | |
| 190 | def download_and_save_text_encoder(output_path: str): |
no test coverage detected
searching dependent graphs…