first adds noise to the input image and keep the added noise for latter rectification.

comfy_extras/nodes_cogvidex_model.py

@@ -1,5 +1,5 @@
 import os
-from typing import Callable, Dict, List, Optional, Tuple, Union
+from typing import List, Optional, Tuple, Union
 import torch
 import folder_paths
 import comfy.model_management as mm
@@ -8,8 +8,7 @@ from diffusers.schedulers import CogVideoXDDIMScheduler, CogVideoXDPMScheduler
 from diffusers.models import AutoencoderKLCogVideoX, CogVideoXTransformer3DModel
 from diffusers import CogVideoXPipeline
 from diffusers.pipelines.cogvideo.pipeline_cogvideox import retrieve_timesteps, get_resize_crop_region_for_grid, get_3d_rotary_pos_embed
-from diffusers.utils.torch_utils import randn_tensor
+from diffusers.utils.torch_utils import randn_tensor 
-from diffusers.callbacks import MultiPipelineCallbacks, PipelineCallback
 
 from diffusers.image_processor import VaeImageProcessor
 from transformers.models.t5.modeling_t5 import T5EncoderModel
@@ -108,8 +107,8 @@ def prepare_latents(
     if latents is None:
         latents = randn_tensor(shape, generator=generator, device=device, dtype=dtype)
     else:
-        if latents.shape != shape:
+        # if latents.shape != shape:
-                raise ValueError(f"Unexpected latents shape, got {latents.shape}, expected {shape}")
+        #         raise ValueError(f"Unexpected latents shape, got {latents.shape}, expected {shape}")
         latents = latents.to(device)
 
     # scale the initial noise by the standard deviation required by the scheduler
@@ -247,21 +246,26 @@ def cogvideox_sampler(
         height // cogvideo_pipeline.vae_scale_factor_spatial,
         width // cogvideo_pipeline.vae_scale_factor_spatial,
     )
- 
+   
-    frames_needed = frames_shape[1]
-    
     # https://arxiv.org/pdf/2403.02827 Our method first adds noise to the input image and keep the added noise for latter rectification.
-    noise = latents.clone()
+
     init_latents = None
-    video_length = (num_frames - 1) // cogvideo_pipeline.vae_scale_factor_temporal + 1
+    noise = latents.clone()
-    if frames_needed > current_frames:
+    video_length = frames_shape[1]
+    if video_length > current_frames:
         # images lanets frame is index 0 ，get the noise latents 
-        init_latents = latents[:, 0, :, :, :, :]
+        init_latents = latents[:, 0:1, :, :, :] 
-    elif frames_needed < current_frames:
+        repeat_factor = video_length - current_frames
-        latents = latents[:, :frames_needed, :, :, :]
+        additional_frame = torch.randn((latents.size(0), repeat_factor, latents.size(2), latents.size(3), latents.size(4)), dtype=latents.dtype, device=latents.device)
-    
+        latents = torch.cat((latents, additional_frame), dim=1)
-    if init_latents is not None:
+    elif video_length < current_frames: 
-        latents = cogvideo_pipeline.scheduler.add_noise(init_latents, noise, latent_timestep)
+        init_latents = latents[:, :video_length, :, :, :]
+        latents = randn_tensor(frames_shape, generator=generator, device=device, dtype=cogvideo_pipeline.vae.dtype)
+    else:
+        init_latents = latents
+        latents = randn_tensor(frames_shape, generator=generator, device=device, dtype=cogvideo_pipeline.vae.dtype)
+
+    latents = cogvideo_pipeline.scheduler.add_noise(init_latents, latents, latent_timestep)
     latents = latents.to(cogvideo_pipeline.transformer.dtype)
 
     # 6. Prepare extra step kwargs. TODO: Logic should ideally just be moved out of the pipeline
@@ -320,7 +324,7 @@ def cogvideox_sampler(
 
                 if i >= len(timesteps) * noise_rectification_period[0] and i < len(timesteps) * noise_rectification_period[1]:
                     delta_frames = noise - noise_pred
-                    delta_noise_adjust = noise_rectification_weight * (delta_frames[:,:,[0],:,:].repeat((1, video_length, 1, 1, 1))) + \
+                    delta_noise_adjust = noise_rectification_weight * (delta_frames[:,[0],:,:,:].repeat((1, video_length, 1, 1, 1))) + \
                                         (1 - noise_rectification_weight) * delta_frames
                     noise_pred = noise_pred + delta_noise_adjust
 
@@ -458,6 +462,37 @@ class CogVideoPipeExtra:
         return (t5_tokenizer, text_encoder, vae3d, transformer, scheduler, dtype)
      
 
+class CogVideoTextEncode:
+    @classmethod
+    def INPUT_TYPES(s):
+        return {"required": {
+            "clip": ("CLIP",),
+            "prompt": ("STRING", {"default": "", "multiline": True, } ),
+            },
+            "optional": {
+                
+                "text": ("STRING", {"forceInput": True}),
+            }
+        }
+
+    RETURN_TYPES = ("CONDITIONING",)
+    RETURN_NAMES = ("conditioning",)
+    FUNCTION = "process"
+    CATEGORY = "CogVideoWrapper"
+
+    def process(self, clip, text, prompt):
+        load_device = mm.text_encoder_device()
+        offload_device = mm.text_encoder_offload_device()
+        clip.tokenizer.t5xxl.pad_to_max_length = True
+        clip.tokenizer.t5xxl.max_length = 226
+        clip.cond_stage_model.to(load_device)
+        tokens = clip.tokenize(prompt or text, return_word_ids=True)
+
+        embeds = clip.encode_from_tokens(tokens, return_pooled=False, return_dict=False)
+        clip.cond_stage_model.to(offload_device)
+
+        return (embeds, )
+    
 
 class CogVideoEncodePrompt:
     @classmethod
@@ -466,8 +501,13 @@ class CogVideoEncodePrompt:
             "t5_tokenizer": ("T5_TOKENIZER",),
             "text_encoder": ("TEXT_ENCODER",),
             "dtype": ("DTYPE",),
-            "prompt": ("STRING", {"default": "", "multiline": True} ),
+            "prompt": ("STRING", {"default": "", "multiline": True, } ),
-            "negative_prompt": ("STRING", {"default": "", "multiline": True} ),
+            "negative_prompt": ("STRING", {"default": "", "multiline": True,}, ),
+            },
+            "optional": {
+                    
+                "negative_text": ("STRING",{"forceInput": True}),
+                "positive_text": ("STRING",{"forceInput": True}),
             }
         }
 
@@ -602,7 +642,7 @@ class CogVideoEncodePrompt:
         return prompt_embeds, negative_prompt_embeds
 
 
-    def process(self, t5_tokenizer, text_encoder, dtype, prompt, negative_prompt):
+    def process(self, t5_tokenizer, text_encoder, dtype, positive_text, negative_text, prompt, negative_prompt):
         device = mm.get_torch_device()
         offload_device = mm.unet_offload_device() 
 
@@ -611,8 +651,8 @@ class CogVideoEncodePrompt:
         positive, negative = self._encode_prompt(
             t5_tokenizer,
             text_encoder,
-            prompt=prompt,
+            prompt=prompt or positive_text,
-            negative_prompt=negative_prompt,
+            negative_prompt=negative_prompt or negative_text,
             do_classifier_free_guidance=True,
             num_videos_per_prompt=1,
             max_sequence_length=226,
@@ -874,6 +914,7 @@ NODE_CLASS_MAPPINGS = {
     "CogVideoModelLoader": CogVideoModelLoader,
     "CogVideoSchedulerLoader": CogVideoSchedulerLoader,
     "CogVideoPipeExtra": CogVideoPipeExtra,
+    "CogVideoTextEncode": CogVideoTextEncode,
     "CogVideoEncodePrompt": CogVideoEncodePrompt,
     "CogVideoSampler": CogVideoSampler,
     "CogVideoImageEncodeSampler": CogVideoImageEncodeSampler,
@@ -884,6 +925,7 @@ NODE_DISPLAY_NAME_MAPPINGS = {
     "CogVideoModelLoader": "CogVideo ModelLoader",
     "CogVideoSchedulerLoader": "CogVideo SchedulerLoader",
     "CogVideoPipeExtra": "CogVideo PipeExtra",
+    "CogVideoTextEncode": "CogVideo TextEncode",
     "CogVideoEncodePrompt": "CogVideo EncodePrompt",
     "CogVideoSampler": "CogVideo Sampler",
     "CogVideoDecode": "CogVideo Decode",