update cli demo with 12GB memory

inference/cli_demo.py

@@ -3,13 +3,13 @@ This script demonstrates how to generate a video from a text prompt using CogVid
 
 Note:
     This script requires the `diffusers>=0.30.0` library to be installed.
-    If the video exported using OpenCV appears “completely green” and cannot be viewed, lease switch to a different player to watch it. This is a normal phenomenon.
 
 Run the script:
     $ python cli_demo.py --prompt "A girl ridding a bike." --model_path THUDM/CogVideoX-2b
 
 """
 
+import gc
 import argparse
 import tempfile
 from typing import Union, List
@@ -18,11 +18,10 @@ import PIL
 import imageio
 import numpy as np
 import torch
-from diffusers import CogVideoXPipeline
+from diffusers import CogVideoXPipeline, CogVideoXDDIMScheduler
-
 
 def export_to_video_imageio(
-    video_frames: Union[List[np.ndarray], List[PIL.Image.Image]], output_video_path: str = None, fps: int = 8
+        video_frames: Union[List[np.ndarray], List[PIL.Image.Image]], output_video_path: str = None, fps: int = 8
 ) -> str:
     """
     Export the video frames to a video file using imageio lib to Avoid "green screen" issue (for example CogVideoX)
@@ -38,14 +37,13 @@ def export_to_video_imageio(
 
 
 def generate_video(
-    prompt: str,
+        prompt: str,
-    model_path: str,
+        model_path: str,
-    output_path: str = "./output.mp4",
+        output_path: str = "./output.mp4",
-    num_inference_steps: int = 50,
+        num_inference_steps: int = 50,
-    guidance_scale: float = 6.0,
+        guidance_scale: float = 6.0,
-    num_videos_per_prompt: int = 1,
+        num_videos_per_prompt: int = 1,
-    device: str = "cuda",
+        dtype: torch.dtype = torch.float16,
-    dtype: torch.dtype = torch.float16,
 ):
     """
     Generates a video based on the given prompt and saves it to the specified path.
@@ -57,36 +55,44 @@ def generate_video(
     - num_inference_steps (int): Number of steps for the inference process. More steps can result in better quality.
     - guidance_scale (float): The scale for classifier-free guidance. Higher values can lead to better alignment with the prompt.
     - num_videos_per_prompt (int): Number of videos to generate per prompt.
-    - device (str): The device to use for computation (e.g., "cuda" or "cpu").
     - dtype (torch.dtype): The data type for computation (default is torch.float16).
+
     """
 
-    # Load the pre-trained CogVideoX pipeline with the specified precision (float16) and move it to the specified device
+    # 1.  Load the pre-trained CogVideoX pipeline with the specified precision (float16).
-    # add device_map="balanced" in the from_pretrained function and remove
+    # add device_map="balanced" in the from_pretrained function and remove the enable_model_cpu_offload()
-    # `pipe.enable_model_cpu_offload()` to enable Multi GPUs (2 or more and each one must have more than 20GB memory) inference.
+    # function to use Multi GPUs.
+
     pipe = CogVideoXPipeline.from_pretrained(model_path, torch_dtype=dtype)
+
+    # 2. Set Scheduler.
+    # Can be changed to `CogVideoXDPMScheduler` or `CogVideoXDDIMScheduler`.
+    # We recommend using `CogVideoXDDIMScheduler` for better results.
+    pipe.scheduler = CogVideoXDDIMScheduler.from_config(pipe.scheduler.config, timestep_spacing="trailing")
+
+    # 3. Enable CPU offload for the model and reset the memory, enable tiling.
     pipe.enable_model_cpu_offload()
 
-    # Encode the prompt to get the prompt embeddings
+    gc.collect()
-    prompt_embeds, _ = pipe.encode_prompt(
+    torch.cuda.empty_cache()
-        prompt=prompt,  # The textual description for video generation
+    torch.cuda.reset_accumulated_memory_stats()
-        negative_prompt=None,  # The negative prompt to guide the video generation
+    torch.cuda.reset_peak_memory_stats()
-        do_classifier_free_guidance=True,  # Whether to use classifier-free guidance
-        num_videos_per_prompt=num_videos_per_prompt,  # Number of videos to generate per prompt
-        max_sequence_length=226,  # Maximum length of the sequence, must be 226
-        device=device,  # Device to use for computation
-        dtype=dtype,  # Data type for computation
-    )
 
-    # Generate the video frames using the pipeline
+    # Using with diffusers branch `cogvideox-followup` to enable tiling. not support in `main` branch.
+    # This will cost ONLY 12GB GPU memory.
+    # pipe.vae.enable_tiling()
+
+    # 4. Generate the video frames based on the prompt.
     video = pipe(
+        prompt=prompt,
+        num_videos_per_prompt=num_videos_per_prompt,  # Number of videos to generate per prompt
         num_inference_steps=num_inference_steps,  # Number of inference steps
+        num_frames=48,  # Number of frames to generate
         guidance_scale=guidance_scale,  # Guidance scale for classifier-free guidance
-        prompt_embeds=prompt_embeds,  # Encoded prompt embeddings
+        generator=torch.Generator().manual_seed(42),  # Set the seed for reproducibility
-        negative_prompt_embeds=torch.zeros_like(prompt_embeds),  # Not Supported negative prompt
     ).frames[0]
 
-    # Export the generated frames to a video file. fps must be 8
+    # 5. Export the generated frames to a video file. fps must be 8
     export_to_video_imageio(video, output_path, fps=8)
 
 
@@ -104,10 +110,6 @@ if __name__ == "__main__":
     )
     parser.add_argument("--guidance_scale", type=float, default=6.0, help="The scale for classifier-free guidance")
     parser.add_argument("--num_videos_per_prompt", type=int, default=1, help="Number of videos to generate per prompt")
-    parser.add_argument(
-        "--device", type=str, default="cuda", help="The device to use for computation (e.g., 'cuda' or 'cpu')"
-    )
-
     parser.add_argument(
         "--dtype", type=str, default="float16", help="The data type for computation (e.g., 'float16' or 'float32')"
     )
@@ -125,6 +127,5 @@ if __name__ == "__main__":
         num_inference_steps=args.num_inference_steps,
         guidance_scale=args.guidance_scale,
         num_videos_per_prompt=args.num_videos_per_prompt,
-        device=args.device,
         dtype=dtype,
     )

requirements.txt

@@ -8,7 +8,7 @@ torchvision==0.19.0
 gradio==4.40.0 # For HF gradio demo
 pillow==9.5.0 # For HF gradio demo
 streamlit==1.37.0 # For streamlit web demo
-opencv-python==4.10 # For diffusers inference origin export video
 imageio==2.34.2 # For diffusers inference export video
 imageio-ffmpeg==0.5.1 # For diffusers inference export video
-openai==1.38.0 # For prompt refiner
+openai==1.40.6 # For prompt refiner
+moviepy==1.0.3 # For export video