Add Diffusers AutoModel support

README.md

@@ -12,14 +12,14 @@ tags:
 
 # 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. 
+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/ | ![compare_autoencoders_1](https://cdn-uploads.huggingface.co/production/uploads/64429aaf7feb866811b12f73/u7ZnjY8FAwu09-iyEC_um.png) |
 | https://www.pexels.com/photo/brown-hummingbird-selective-focus-photography-1133957/ | ![compare_autoencoders_2](https://cdn-uploads.huggingface.co/production/uploads/64429aaf7feb866811b12f73/ZzvJu3VfrzlvZ7bDDASog.png) |
 | https://www.pexels.com/photo/person-with-body-painting-1209843/ | ![compare_autoencoders_3](https://cdn-uploads.huggingface.co/production/uploads/64429aaf7feb866811b12f73/B56LPhLYiGT0ffnBVIRbP.png) |
@@ -51,4 +51,20 @@ with torch.inference_mode():
 
 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)
+```

config.json

@@ -1,6 +1,9 @@
 {
   "_class_name": "Flux2TinyAutoEncoder",
   "_diffusers_version": "0.35.2",
+  "auto_map": {
+    "AutoModel": "flux2_tiny_autoencoder.Flux2TinyAutoEncoder"
+  },
   "act_fn": "silu",
   "decoder_block_out_channels": [
     64,