Update README.md

README.md

@@ -32,20 +32,21 @@ 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",
-    device="cuda"
-)
-tiny_vae.eval()
+).to(device=device, dtype=torch.bfloat16)
 
-pil_image = Image.open("/path/to/image.png")
+pil_image = Image.open("pexels.jpg")
 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")