Update README.md

README.md

@@ -31,16 +31,14 @@ import torch
 from torchvision import transforms
 
 def preview_image(latents, pipe):
-    tea = TeaDecoder(ch_in=16)
-    load_model(tea, './vae_decoder.safetensors')
-    tea.eval()
     latents = FluxPipeline._unpack_latents(latents,
                                            pipe.default_sample_size * pipe.vae_scale_factor,
                                            pipe.default_sample_size * pipe.vae_scale_factor,
                                            pipe.vae_scale_factor)
+    tea = TeaDecoder(ch_in=16)
+    load_model(tea, './vae_decoder.safetensors')
     tea = tea.to(device='cuda')
     output = tea(latents.to(torch.float32)) / 2.0 + 0.5
-    tea = tea.to(device='cpu')
     preview = transforms.ToPILImage()(output[0].clamp(0, 1))
 
     return preview
@@ -65,7 +63,7 @@ if __name__ == '__main__':
     # Return the upscaled and preview image.
     upscaled = full_size_image(latents, pipe)
     preview = preview_image(latents, pipe)
-    preview.save('cat.preview.png')
+    preview.save('cat.png')
 ```
 
 ## Disclaimer