Update nodes_rf_inversion.py

kijai · Dec 23, 2024 · 46e31f1 · 46e31f1
1 parent ebb267b
commit 46e31f1
Showing 1 changed file with 12 additions and 0 deletions.
diff --git a/nodes_rf_inversion.py b/nodes_rf_inversion.py
@@ -133,6 +133,10 @@ def process(self, model, hyvid_embeds, flow_shift, steps, embedded_guidance_scal
                 offload_txt_in = model["block_swap_args"]["offload_txt_in"],
                 offload_img_in = model["block_swap_args"]["offload_img_in"],
             )
+        elif model["auto_cpu_offload"]:
+            for name, param in transformer.named_parameters():
+                if "single" not in name and "double" not in name:
+                    param.data = param.data.to(device)
         elif model["manual_offloading"]:
             transformer.to(device)
 
@@ -344,6 +348,10 @@ def process(self, model, hyvid_embeds, flow_shift, steps, embedded_guidance_scal
                 offload_txt_in = model["block_swap_args"]["offload_txt_in"],
                 offload_img_in = model["block_swap_args"]["offload_img_in"],
             )
+        elif model["auto_cpu_offload"]:
+            for name, param in transformer.named_parameters():
+                if "single" not in name and "double" not in name:
+                    param.data = param.data.to(device)
         elif model["manual_offloading"]:
             transformer.to(device)
 
@@ -539,6 +547,10 @@ def process(self, model, width, height, num_frames, hyvid_embeds, hyvid_embeds_2
                 offload_txt_in = model["block_swap_args"]["offload_txt_in"],
                 offload_img_in = model["block_swap_args"]["offload_img_in"],
             )
+        elif model["auto_cpu_offload"]:
+            for name, param in transformer.named_parameters():
+                if "single" not in name and "double" not in name:
+                    param.data = param.data.to(device)
         elif model["manual_offloading"]:
             transformer.to(device)