#!/usr/bin/env python

import os
import sys
import torch
import shutil
from diffusers import StableDiffusionPipeline, AutoencoderKL

# append project directory to path so predict.py can be imported
sys.path.append('.')

from predict import MODEL_NAME, MODEL_CACHE, VAE_NAME, VAE_CACHE

if os.path.exists(MODEL_CACHE):
    shutil.rmtree(MODEL_CACHE)
os.makedirs(MODEL_CACHE, exist_ok=True)

if os.path.exists(VAE_CACHE):
    shutil.rmtree(VAE_CACHE)
os.makedirs(VAE_CACHE, exist_ok=True)

url = "https://huggingface.co/stabilityai/sd-vae-ft-mse-original/blob/main/vae-ft-mse-840000-ema-pruned.safetensors"
vae = AutoencoderKL.from_single_file(
    url,
    cache_dir=VAE_CACHE
)

pipe = StableDiffusionPipeline.from_pretrained(
    MODEL_NAME,
    torch_dtype=torch.float16,
    use_safetensors=True,
)
pipe.save_pretrained(MODEL_CACHE, safe_serialization=True)