Spaces:

HTBris
/

zerogpu-sdxlTrainer

Running on Zero

zerogpu-sdxlTrainer / toolkit /util /vae.py

Deploy Ostris AI-Toolkit ZeroGPU Trainer (Illustrious XL / SDXL / FLUX.1)

855e1f5 3 days ago

520 Bytes

from diffusers import AutoencoderKL


def load_vae(vae_path, dtype):
    try:
        vae = AutoencoderKL.from_pretrained(
            vae_path,
            torch_dtype=dtype,
        )
    except Exception as e:
        try:
            vae = AutoencoderKL.from_pretrained(
                vae_path,
                subfolder="vae",
                torch_dtype=dtype,
            )
        except Exception as e:
            raise ValueError(f"Failed to load VAE from {vae_path}: {e}")
    vae.to(dtype)
    return vae

	from diffusers import AutoencoderKL


	def load_vae(vae_path, dtype):
	try:
	vae = AutoencoderKL.from_pretrained(
	vae_path,
	torch_dtype=dtype,
	)
	except Exception as e:
	try:
	vae = AutoencoderKL.from_pretrained(
	vae_path,
	subfolder="vae",
	torch_dtype=dtype,
	)
	except Exception as e:
	raise ValueError(f"Failed to load VAE from {vae_path}: {e}")
	vae.to(dtype)
	return vae