| --- |
| library_name: diffusers |
| license: apache-2.0 |
| datasets: |
| - laion/relaion400m |
| base_model: |
| - black-forest-labs/FLUX.2-dev |
| tags: |
| - tae |
| - taef2 |
| --- |
| |
| # About |
|
|
| Tiny AutoEncoder trained on the latent space of [black-forest-labs/FLUX.2-dev](https://huggingface.co/black-forest-labs/FLUX.2-dev)'s autoencoder. Works to convert between latent and image space up to 20x faster and in 28x fewer parameters at the expense of a small amount of quality. |
|
|
| Code for this model is available [here](https://huggingface.co/fal/FLUX.2-Tiny-AutoEncoder/blob/main/flux2_tiny_autoencoder.py). |
|
|
| # Round-Trip Comparisons |
|
|
| | Source | Image | |
| | ------ | ----- | |
| | https://www.pexels.com/photo/mirror-lying-on-open-book-11495792/ |  | |
| | https://www.pexels.com/photo/brown-hummingbird-selective-focus-photography-1133957/ |  | |
| | https://www.pexels.com/photo/person-with-body-painting-1209843/ |  | |
|
|
| # Example Usage |
|
|
| ```py |
| import torch |
| import torchvision.transforms.functional as F |
| |
| from PIL import Image |
| from flux2_tiny_autoencoder import Flux2TinyAutoEncoder |
| |
| device = torch.device("cuda") |
| tiny_vae = Flux2TinyAutoEncoder.from_pretrained( |
| "fal/FLUX.2-Tiny-AutoEncoder", |
| ).to(device=device, dtype=torch.bfloat16) |
| |
| pil_image = Image.open("/path/to/image.png") |
| image_tensor = F.to_tensor(pil_image) |
| image_tensor = image_tensor.unsqueeze(0) * 2.0 - 1.0 |
| image_tensor = image_tensor.to(device, dtype=tiny_vae.dtype) |
| |
| with torch.inference_mode(): |
| latents = tiny_vae.encode(image_tensor, return_dict=False) |
| recon = tiny_vae.decode(latents, return_dict=False) |
| recon = recon.squeeze(0).clamp(-1, 1) / 2.0 + 0.5 |
| recon = recon.float().detach().cpu() |
| |
| recon_image = F.to_pil_image(recon) |
| recon_image.save("reconstituted.png") |
| ``` |
|
|
| ## Use with Diffusers 🧨 |
|
|
| ```py |
| import torch |
| from diffusers import AutoModel, Flux2Pipeline |
| |
| device = torch.device("cuda") |
| tiny_vae = AutoModel.from_pretrained( |
| "fal/FLUX.2-Tiny-AutoEncoder", trust_remote_code=True, torch_dtype=torch.bfloat16 |
| ).to(device) |
| |
| pipe = Flux2Pipeline.from_pretrained( |
| "black-forest-labs/FLUX.2-dev", vae=tiny_vae, torch_dtype=torch.bfloat16 |
| ).to(device) |
| ``` |
|
|
