From 21675b296b761ad9fe8f8be1b02606110fb69399 Mon Sep 17 00:00:00 2001
From: kijai <40791699+kijai@users.noreply.github.com>
Date: Sun, 22 Sep 2024 16:07:32 +0300
Subject: [PATCH] Add 5b I2V GGUF

https://huggingface.co/Kijai/CogVideoX_GGUF/blob/main/CogVideoX_5b_I2V_GGUF_Q4_0.safetensors
---
 cogvideox_fun/pipeline_cogvideox_inpaint.py |  1 -
 configs/transformer_config_I2V_5b.json      |  2 --
 nodes.py                                    | 12 +++++++++++-
 3 files changed, 11 insertions(+), 4 deletions(-)

diff --git a/cogvideox_fun/pipeline_cogvideox_inpaint.py b/cogvideox_fun/pipeline_cogvideox_inpaint.py
index f372342..7437d03 100644
--- a/cogvideox_fun/pipeline_cogvideox_inpaint.py
+++ b/cogvideox_fun/pipeline_cogvideox_inpaint.py
@@ -21,7 +21,6 @@ from typing import Callable, Dict, List, Optional, Tuple, Union
 import torch
 import torch.nn.functional as F
 from einops import rearrange
-from transformers import T5EncoderModel, T5Tokenizer
 
 from diffusers.callbacks import MultiPipelineCallbacks, PipelineCallback
 from diffusers.models import AutoencoderKLCogVideoX, CogVideoXTransformer3DModel
diff --git a/configs/transformer_config_I2V_5b.json b/configs/transformer_config_I2V_5b.json
index 3265e76..420eb9f 100644
--- a/configs/transformer_config_I2V_5b.json
+++ b/configs/transformer_config_I2V_5b.json
@@ -1,6 +1,4 @@
 {
-    "_class_name": "CogVideoXTransformer3DModel",
-    "_diffusers_version": "0.31.0.dev0",
     "activation_fn": "gelu-approximate",
     "attention_bias": true,
     "attention_head_dim": 64,
diff --git a/nodes.py b/nodes.py
index 4577310..5be183e 100644
--- a/nodes.py
+++ b/nodes.py
@@ -173,6 +173,7 @@ class DownloadAndLoadCogVideoGGUFModel:
                 "model": (
                     [
                         "CogVideoX_5b_GGUF_Q4_0.safetensors",
+                        "CogVideoX_5b_I2V_GGUF_Q4_0.safetensors",
                         "CogVideoX_5b_fun_GGUF_Q4_0.safetensors",
                     ],
                 ),
@@ -198,11 +199,15 @@ class DownloadAndLoadCogVideoGGUFModel:
         if not os.path.exists(gguf_path):
             gguf_path = os.path.join(download_path, model)
             if not os.path.exists(gguf_path):
+                if "I2V" in model:
+                    repo_id = "Kijai/CogVideoX_GGUF"
+                else:
+                    repo_id = "MinusZoneAI/ComfyUI-CogVideoX-MZ"
                 log.info(f"Downloading model to: {gguf_path}")
                 from huggingface_hub import snapshot_download
 
                 snapshot_download(
-                    repo_id="MinusZoneAI/ComfyUI-CogVideoX-MZ",
+                    repo_id=repo_id,
                     allow_patterns=[f"*{model}*"],
                     local_dir=download_path,
                     local_dir_use_symlinks=False,
@@ -212,6 +217,8 @@ class DownloadAndLoadCogVideoGGUFModel:
         with open(os.path.join(script_directory, 'configs', 'transformer_config_5b.json')) as f:
                     transformer_config = json.load(f)
         sd = load_torch_file(gguf_path)
+        for key, value in sd.items():
+            print(key, value.shape, value.dtype)
 
         from . import mz_gguf_loader
         import importlib
@@ -221,6 +228,9 @@ class DownloadAndLoadCogVideoGGUFModel:
             if "fun" in model:
                 transformer_config["in_channels"] = 33
                 transformer = CogVideoXTransformer3DModelFun.from_config(transformer_config)
+            elif "I2V" in model:
+                transformer_config["in_channels"] = 32
+                transformer = CogVideoXTransformer3DModel.from_config(transformer_config)
             else:
                 transformer_config["in_channels"] = 16
                 transformer = CogVideoXTransformer3DModel.from_config(transformer_config)