INNER CODE UNIT · Python

tinyVae

kiri-art/docker-diffusers-api · api/app.py:73

def tinyVae(origVae: AutoencoderKL):
    global tiny_vae
    if not tiny_vae:
        tiny_vae = AutoencoderTiny.from_pretrained(
            "madebyollin/taesd",
            torch_dtype=torch.float16,
            in_channels=origVae.config.in_channels,
            out_channels=origVae.config.out_channels,
            act_fn=origVae.config.act_fn,
            latent_channels=origVae.config.latent_channels,
            scaling_factor=origVae.config.scaling_factor,
            force_upcast=origVae.config.force_upcast,
        )
        tiny_vae.to("cuda")

    return tiny_vae

View source record →