From 87ed4a56cf48cf8d71e07f3d15088d2f2d196538 Mon Sep 17 00:00:00 2001
From: kijai <40791699+kijai@users.noreply.github.com>
Date: Sun, 10 Nov 2024 18:19:35 +0200
Subject: [PATCH] Update model_loading.py

---
 model_loading.py | 2 +-
 1 file changed, 1 insertion(+), 1 deletion(-)

diff --git a/model_loading.py b/model_loading.py
index 08218ca..acf8c9e 100644
--- a/model_loading.py
+++ b/model_loading.py
@@ -418,8 +418,8 @@ class DownloadAndLoadCogVideoGGUFModel:
                 else:
                     transformer = CogVideoXTransformer3DModel.from_config(transformer_config)
 
+            params_to_keep = {"patch_embed", "pos_embedding", "time_embedding"}
             if "2b" in model:
-                params_to_keep = {"patch_embed", "pos_embedding", "time_embedding"}
                 cast_dtype = torch.float16
             elif "1_5" in model:
                 params_to_keep = {"norm1.linear.weight", "patch_embed", "time_embedding", "ofs_embedding", "norm_final", "norm_out", "proj_out"}