TAEF1: Tiny Autoencoder for FLUX.1

Introduction

TAEF1 is a compact autoencoder designed to operate with the "latent API" used by FLUX.1's Variational Autoencoder (VAE). It is intended for real-time previewing of FLUX.1's generation process. The repository includes .safetensors versions of the TAEF1 weights.

Architecture

TAEF1 is structured to integrate seamlessly with the FLUX.1 pipeline, using a minimalistic architecture to achieve quick and efficient encoding and decoding. Its small size allows it to support real-time applications effectively.

Training

The model leverages the latent API to optimize for performance in generating previews. It is designed to be lightweight, ensuring fast inference times suitable for applications requiring immediate visual feedback.

Guide: Running Locally

To use TAEF1 in conjunction with the Diffusers library, follow these steps:

Install Pre-requisites: Ensure you have PyTorch and the Diffusers library installed.

Load the Model:

import torch
from diffusers import FluxPipeline, AutoencoderTiny

pipe = FluxPipeline.from_pretrained(
    "black-forest-labs/FLUX.1-schnell", torch_dtype=torch.bfloat16
)
pipe.vae = AutoencoderTiny.from_pretrained("madebyollin/taef1", torch_dtype=torch.bfloat16)
pipe.enable_sequential_cpu_offload()

Generate an Image:

prompt = "slice of delicious New York-style berry cheesecake"
image = pipe(
    prompt,
    guidance_scale=0.0,
    num_inference_steps=4,
    max_sequence_length=256,
).images[0]
image.save("cheesecake.png")

Hardware Recommendation: For optimal performance, using a cloud GPU such as those provided by AWS, Google Cloud, or Azure is recommended.

License

This project is released under the MIT License, allowing for open use and modification.

More Related APIs