From f3b365a5dabf836374034250f877fb0f15af3812 Mon Sep 17 00:00:00 2001
From: DeepBeepMeep <deepbeepmeep@yahoo.com>
Date: Sun, 2 Mar 2025 17:55:30 +0100
Subject: [PATCH] RIFLEx support

---
 wan/modules/vae.py | 20 +++++++++-----------
 1 file changed, 9 insertions(+), 11 deletions(-)

diff --git a/wan/modules/vae.py b/wan/modules/vae.py
index a4a68d2..ded9f92 100644
--- a/wan/modules/vae.py
+++ b/wan/modules/vae.py
@@ -677,16 +677,14 @@ class WanVAE:
         """
         videos: A list of videos each with shape [C, T, H, W].
         """
-        with amp.autocast(dtype=self.dtype):
-            return [
-                self.model.encode(u.unsqueeze(0), self.scale).float().squeeze(0)
-                for u in videos
-            ]
+        return [
+            self.model.encode(u.unsqueeze(0), self.scale).float().squeeze(0)
+            for u in videos
+        ]
 
     def decode(self, zs):
-        with amp.autocast(dtype=self.dtype):
-            return [
-                self.model.decode(u.unsqueeze(0),
-                                  self.scale).float().clamp_(-1, 1).squeeze(0)
-                for u in zs
-            ]
+        return [
+            self.model.decode(u.unsqueeze(0),
+                                self.scale).float().clamp_(-1, 1).squeeze(0)
+            for u in zs
+        ]