From c6eca73aaa185df2b85f597fc30580a296dce3e9 Mon Sep 17 00:00:00 2001
From: wailovet <wailovet@163.com>
Date: Sun, 1 Sep 2024 03:29:35 +0800
Subject: [PATCH] 'update'

---
 __init__.py          | 2 +-
 mz_cogvideox_core.py | 7 +++++--
 2 files changed, 6 insertions(+), 3 deletions(-)

diff --git a/__init__.py b/__init__.py
index ea135b2..94d3fdb 100644
--- a/__init__.py
+++ b/__init__.py
@@ -28,7 +28,7 @@ class MZ_CogVideoXLoader:
                 "vae_name": (folder_paths.get_filename_list("vae"), ),
                 "weight_dtype": (["default", "fp8_e4m3fn", "fp8_e5m2"],),
                 "fp8_fast_mode": ("BOOLEAN", {"default": False}),
-                "dyn_offload_cpu_layer": ("INT", {"default": 0, "tooltip": "0-42,默认0不启用"}),
+                # "dyn_offload_cpu_layer": ("INT", {"default": 0, "tooltip": "0-42,默认0不启用"}),
                 "enable_sequential_cpu_offload": ("BOOLEAN", {"default": False, "tooltip": "significantly reducing memory usage and slows down the inference"}),
             }}
 
diff --git a/mz_cogvideox_core.py b/mz_cogvideox_core.py
index e013a26..227dd16 100644
--- a/mz_cogvideox_core.py
+++ b/mz_cogvideox_core.py
@@ -269,14 +269,17 @@ def MZ_CogVideoXLoader_call(args={}):
 
     transformer.load_state_dict(unet_sd)
 
-    dtype = torch.bfloat16
+    dtype = None
     weight_dtype = args.get("weight_dtype")
     if weight_dtype == "fp8_e4m3fn":
         dtype = torch.float8_e4m3fn
+        transformer.to(dtype)
     elif weight_dtype == "fp8_e5m2":
         dtype = torch.float8_e5m2
+        transformer.to(dtype)
+    else:
+        dtype = transformer.parameters().__next__().dtype
 
-    transformer.to(dtype)
     if dtype == torch.float8_e4m3fn or dtype == torch.float8_e5m2:
         fp8_fast_mode = args.get("fp8_fast_mode", False)
         if fp8_fast_mode: