Update README.md

README.md

@@ -23,23 +23,29 @@ Code for this model is available [here](https://huggingface.co/fal/FLUX.2-Tiny-A
 | 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) |
 
-# Usage
+# Example Usage
 
 ```py
 import torch
-import torchvision.transformers.functional as F
+import torchvision.transforms.functional as F
 
 from PIL import Image
 from flux2_tiny_autoencoder import Flux2TinyAutoEncoder
 
-tiny_vae = Flux2TinyAutoEncoder.from_pretrained("fal/FLUX.2-Tiny-AutoEncoder")
+tiny_vae = Flux2TinyAutoEncoder.from_pretrained(
+    "fal/FLUX.2-Tiny-AutoEncoder",
+    device="cuda"
+)
 tiny_vae.eval()
 
-pil_image = Image.open("/path/to/your/image.png")
+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
 
 with torch.inference_mode():
-    latents = tiny_vae.encode(F.to_tensor(pil_image))
-    recon = tiny_vae.decode(latents)
+    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_image = F.to_pil_image(recon)
 recon_image.save("reconstituted.png")