""" This script demonstrates how to generate a video using the CogVideoX model with the Hugging Face `diffusers` pipeline. The script supports different types of video generation, including text-to-video (t2v), image-to-video (i2v), and video-to-video (v2v), depending on the input data and different weight. - text-to-video: THUDM/CogVideoX-5b, THUDM/CogVideoX-2b or THUDM/CogVideoX1.5-5b - video-to-video: THUDM/CogVideoX-5b, THUDM/CogVideoX-2b or THUDM/CogVideoX1.5-5b - image-to-video: THUDM/CogVideoX-5b-I2V or THUDM/CogVideoX1.5-5b-I2V Running the Script: To run the script, use the following command with appropriate arguments: ```bash $ python cli_demo.py --prompt "A girl riding a bike." --model_path THUDM/CogVideoX1.5-5b --generate_type "t2v" ``` You can change `pipe.enable_sequential_cpu_offload()` to `pipe.enable_model_cpu_offload()` to speed up inference, but this will use more GPU memory Additional options are available to specify the model path, guidance scale, number of inference steps, video generation type, and output paths. """ from typing import TYPE_CHECKING, Any, Dict, List, Tuple import argparse import logging import os import sys from typing import Literal, Optional from pathlib import Path import json from datetime import timedelta import random from safetensors.torch import load_file, save_file from tqdm import tqdm from einops import rearrange, repeat import math import numpy as np from PIL import Image import torch from diffusers import ( CogVideoXDPMScheduler, CogVideoXImageToVideoPipeline, CogVideoXPipeline, CogVideoXVideoToVideoPipeline, AutoencoderKLCogVideoX, CogVideoXTransformer3DModel, ) from diffusers.utils import export_to_video, load_image, load_video from peft import LoraConfig, get_peft_model_state_dict, set_peft_model_state_dict sys.path.append(os.path.abspath(os.path.join(sys.path[0], "../"))) from finetune.pipeline.flovd_FVSM_cogvideox_controlnet_pipeline import FloVDCogVideoXControlnetImageToVideoPipeline from finetune.pipeline.flovd_OMSM_cogvideox_pipeline import FloVDOMSMCogVideoXImageToVideoPipeline from finetune.schemas import Components, Args from finetune.modules.cogvideox_controlnet import CogVideoXControlnet from finetune.modules.cogvideox_custom_model import CustomCogVideoXTransformer3DModel from transformers import AutoTokenizer, T5EncoderModel from finetune.modules.camera_sampler import SampleManualCam from finetune.modules.camera_flow_generator import CameraFlowGenerator from finetune.modules.utils import get_camera_flow_generator_input, forward_bilinear_splatting, flow_to_color from finetune.modules.depth_warping.depth_warping import unnormalize_intrinsic from finetune.datasets.utils import ( preprocess_image_with_resize, preprocess_video_with_resize, ) from torch.utils.data import Dataset from torchvision import transforms import torch.distributed as dist from torch.utils.data.distributed import DistributedSampler import pdb sys.path.append(os.path.abspath(os.path.join(sys.path[-1], 'finetune'))) # for camera flow generator os.environ["TOKENIZERS_PARALLELISM"] = "false" logging.basicConfig(level=logging.INFO) # Recommended resolution for each model (width, height) RESOLUTION_MAP = { # cogvideox1.5-* "cogvideox1.5-5b-i2v": (768, 1360), "cogvideox1.5-5b": (768, 1360), # cogvideox-* "cogvideox-5b-i2v": (480, 720), "cogvideox-5b": (480, 720), "cogvideox-2b": (480, 720), } def init_dist(launcher="slurm", backend='nccl', port=29500, **kwargs): """Initializes distributed environment.""" if launcher == 'pytorch': rank = int(os.environ['RANK']) num_gpus = torch.cuda.device_count() local_rank = rank % num_gpus torch.cuda.set_device(local_rank) dist.init_process_group(backend=backend, timeout=timedelta(minutes=30), **kwargs) elif launcher == 'slurm': proc_id = int(os.environ['SLURM_PROCID']) ntasks = int(os.environ['SLURM_NTASKS']) node_list = os.environ['SLURM_NODELIST'] num_gpus = torch.cuda.device_count() local_rank = proc_id % num_gpus torch.cuda.set_device(local_rank) addr = subprocess.getoutput( f'scontrol show hostname {node_list} | head -n1') os.environ['MASTER_ADDR'] = addr os.environ['WORLD_SIZE'] = str(ntasks) os.environ['RANK'] = str(proc_id) port = os.environ.get('PORT', port) os.environ['MASTER_PORT'] = str(port) dist.init_process_group(backend=backend, timeout=timedelta(minutes=30)) else: raise NotImplementedError(f'Not implemented launcher type: `{launcher}`!') # https://github.com/pytorch/pytorch/issues/98763 # torch.cuda.set_device(local_rank) return local_rank def load_cogvideox_flovd_FVSM_controlnet_pipeline(controlnet_path, backbone_path, device, dtype): controlnet_sd = torch.load(controlnet_path, map_location='cpu')['module'] tokenizer = AutoTokenizer.from_pretrained(backbone_path, subfolder="tokenizer") text_encoder = T5EncoderModel.from_pretrained(backbone_path, subfolder="text_encoder") transformer = CustomCogVideoXTransformer3DModel.from_pretrained(backbone_path, subfolder="transformer") vae = AutoencoderKLCogVideoX.from_pretrained(backbone_path, subfolder="vae") scheduler = CogVideoXDPMScheduler.from_pretrained(backbone_path, subfolder="scheduler") additional_kwargs = { 'num_layers': 6, 'out_proj_dim_factor': 64, 'out_proj_dim_zero_init': True, 'notextinflow': True, } controlnet = CogVideoXControlnet.from_pretrained(backbone_path, subfolder="transformer", **additional_kwargs) controlnet.eval() missing, unexpected = controlnet.load_state_dict(controlnet_sd) if len(missing) != 0 or len(unexpected) != 0: print(f"Missing keys : {missing}") print(f"Unexpected keys : {unexpected}") pipe = FloVDCogVideoXControlnetImageToVideoPipeline( tokenizer=tokenizer, text_encoder=text_encoder, vae=vae, transformer=transformer, controlnet=controlnet, scheduler=scheduler, ) # pipe.enable_model_cpu_offload(device=device) pipe = pipe.to(device, dtype) return pipe def load_cogvideox_flovd_OMSM_lora_pipeline(omsm_path, backbone_path, transformer_lora_config, device, dtype): tokenizer = AutoTokenizer.from_pretrained(backbone_path, subfolder="tokenizer") text_encoder = T5EncoderModel.from_pretrained(backbone_path, subfolder="text_encoder") transformer = CogVideoXTransformer3DModel.from_pretrained(backbone_path, subfolder="transformer") vae = AutoencoderKLCogVideoX.from_pretrained(backbone_path, subfolder="vae") scheduler = CogVideoXDPMScheduler.from_pretrained(backbone_path, subfolder="scheduler") # 1) Load Lora weight transformer.add_adapter(transformer_lora_config) lora_state_dict = FloVDOMSMCogVideoXImageToVideoPipeline.lora_state_dict(omsm_path) transformer_state_dict = { f'{k.replace("transformer.", "")}': v for k, v in lora_state_dict.items() if k.startswith("transformer.") } incompatible_keys = set_peft_model_state_dict(transformer, transformer_state_dict, adapter_name="default") if incompatible_keys is not None: # check only for unexpected keys unexpected_keys = getattr(incompatible_keys, "unexpected_keys", None) if unexpected_keys: logger.warning( f"Loading adapter weights from state_dict led to unexpected keys not found in the model: " f" {unexpected_keys}. " ) # 2) Load Other weight load_path = os.path.join(omsm_path, "selected_blocks.safetensors") if os.path.exists(load_path): tensor_dict = load_file(load_path) block_state_dicts = {} for k, v in tensor_dict.items(): block_name, param_name = k.split(".", 1) if block_name not in block_state_dicts: block_state_dicts[block_name] = {} block_state_dicts[block_name][param_name] = v for block_name, state_dict in block_state_dicts.items(): if hasattr(transformer, block_name): getattr(transformer, block_name).load_state_dict(state_dict) else: raise ValueError(f"Transformer has no attribute '{block_name}'") pipe = FloVDOMSMCogVideoXImageToVideoPipeline( tokenizer=tokenizer, text_encoder=text_encoder, vae=vae, transformer=transformer, scheduler=scheduler, ) # pipe.load_lora_weights(omsm_path, weight_name="pytorch_lora_weights.safetensors", adapter_name="test_1") # pipe.fuse_lora(components=["transformer"], lora_scale=1.0) # pipe.enable_model_cpu_offload(device=device) pipe = pipe.to(device, dtype) return pipe class I2VFlowDataset_Inference(Dataset): def __init__( self, max_num_frames: int, height: int, width: int, data_root: str, max_num_videos: int = None, ) -> None: self.train_resolution = (int(max_num_frames), int(height), int(width)) data_root = Path(data_root) metadata_path = data_root / "metadata_reformat.jsonl" assert metadata_path.is_file(), "For this dataset type, you need metadata.jsonl in the root path" metadata = [] with open(metadata_path, "r") as f: for line in f: metadata.append( json.loads(line) ) metadata = random.sample(metadata, max_num_videos) self.prompts = [x["prompt"] for x in metadata] self.prompt_embeddings = [data_root / "prompt_embeddings_revised" / (x["hash_code"] + '.safetensors') for x in metadata] self.videos = [data_root / "video_latent" / "x".join(str(x) for x in self.train_resolution) / (x["hash_code"] + '.safetensors') for x in metadata] self.images = [data_root / "first_frames" / (x["hash_code"] + '.png') for x in metadata] self.flows = [data_root / "flow_direct_f_latent" / (x["hash_code"] + '.safetensors') for x in metadata] self.masks = [data_root / "valid_mask" / (x["hash_code"] + '.bin') for x in metadata] self.max_num_frames = max_num_frames self.height = height self.width = width self.__frame_transforms = transforms.Compose([transforms.Lambda(lambda x: x / 255.0 * 2.0 - 1.0)]) self.__image_transforms = self.__frame_transforms self.length = len(self.videos) print(f"Dataset size: {self.length}") def __len__(self) -> int: return self.length def load_data_pair(self, index): prompt_embedding_path = self.prompt_embeddings[index] encoded_video_path = self.videos[index] encoded_flow_path = self.flows[index] prompt_embedding = load_file(prompt_embedding_path)["prompt_embedding"] encoded_video = load_file(encoded_video_path)["encoded_video"] # CFHW encoded_flow = load_file(encoded_flow_path)["encoded_flow_f"] # CFHW return prompt_embedding, encoded_video, encoded_flow def __getitem__(self, index: int) -> Dict[str, Any]: while True: try: prompt_embedding, encoded_video, encoded_flow = self.load_data_pair(index) break except Exception as e: print(f"Error loading {self.prompt_embeddings[index]}: {str(e)}") index = random.randint(0, self.length - 1) image_path = self.images[index] prompt = self.prompts[index] _, image = self.preprocess(None, image_path) image = self.image_transform(image) # shape of encoded_video: [C, F, H, W] # shape and scale of image: [C, H, W], [-1,1] return { "image": image, "prompt": prompt, "prompt_embedding": prompt_embedding, "encoded_video": encoded_video, "encoded_flow": encoded_flow, "video_metadata": { "num_frames": encoded_video.shape[1], "height": encoded_video.shape[2], "width": encoded_video.shape[3], }, } def preprocess(self, video_path: Path | None, image_path: Path | None) -> Tuple[torch.Tensor, torch.Tensor]: if video_path is not None: video = preprocess_video_with_resize(video_path, self.max_num_frames, self.height, self.width) else: video = None if image_path is not None: image = preprocess_image_with_resize(image_path, self.height, self.width) else: image = None return video, image def video_transform(self, frames: torch.Tensor) -> torch.Tensor: return torch.stack([self.__frame_transforms(f) for f in frames], dim=0) def image_transform(self, image: torch.Tensor) -> torch.Tensor: return self.__image_transforms(image) def initialize_flow_generator(target): depth_estimator_kwargs = { "target": target, "kwargs": { "ckpt_path": '/workspace/workspace/checkpoints/depth_anything/depth_anything_v2_metric_hypersim_vitb.pth', "model_config": { "max_depth": 20, "encoder": 'vitb', "features": 128, "out_channels": [96, 192, 384, 768], } } } return CameraFlowGenerator(depth_estimator_kwargs) def integrate_flow(camera_flow, object_flow, depth_ctxt, camera_flow_generator, camera_flow_generator_input): # camera_flow: (BF)CHW # object_flow: (BF)CHW # depth_ctxt: B1HW B, F = camera_flow_generator_input["target"]["intrinsics"].shape[:2] H, W = object_flow.shape[-2:] c2w_ctxt = repeat(camera_flow_generator_input["context"]["extrinsics"], "b t h w -> (b v t) h w", v=F) # No need to apply inverse as it is an eye matrix. c2w_trgt = rearrange(torch.inverse(camera_flow_generator_input["target"]["extrinsics"]), "b t h w -> (b t) h w") intrinsics_ctxt = unnormalize_intrinsic(repeat(camera_flow_generator_input["context"]["intrinsics"], "b t h w -> (b v t) h w", v=F), size=(H, W)) with torch.cuda.amp.autocast(enabled=False): warped_object_flow = camera_flow_generator.depth_warping_module.warper.forward_warp_displacement( depth1=repeat(depth_ctxt, "b c h w -> (b f) c h w", f=F), flow1=object_flow, transformation1=c2w_ctxt, transformation2=c2w_trgt, intrinsic1=intrinsics_ctxt, intrinsic2=None, ) integrated_flow = camera_flow + warped_object_flow return integrated_flow def save_flow(flow, filename, fps=16): # flow: (BF)CHW, arbitrary scale flow_RGB = flow_to_color(flow) # BF,C,H,W (B=1) frame_list = [] for frame in flow_RGB: frame = (frame.permute(1,2,0).float().detach().cpu().numpy()).astype(np.uint8).clip(0,255) frame_list.append(Image.fromarray(frame)) export_to_video(frame_list, filename, fps=fps) def save_flow_warped_video(image, flow, filename, fps=16): # image: CHW, 0~255 scale # flow: (BF)CHW, arbitrary scale warped_video = forward_bilinear_splatting(repeat(image, 'c h w -> f c h w', f=flow.size(0)), flow.to(torch.float)) frame_list = [] for frame in warped_video: frame = (frame.permute(1,2,0).float().detach().cpu().numpy()).astype(np.uint8).clip(0,255) frame_list.append(Image.fromarray(frame)) export_to_video(frame_list, filename, fps=fps) def generate_video( # prompt: str, launcher: str, port: int, data_root: str, fvsm_path: str, omsm_path: str, num_frames: int = 81, width: Optional[int] = None, height: Optional[int] = None, output_path: str = "./output.mp4", image_path: str = "", num_inference_steps: int = 50, guidance_scale: float = 6.0, num_videos_per_prompt: int = 1, dtype: torch.dtype = torch.bfloat16, seed: int = 42, fps: int = 16, controlnet_guidance_end: float = 0.4, max_num_videos: int = None, use_dynamic_cfg: bool = False, pose_type: str = "manual", speed: float = 0.5, use_flow_integration: bool = False, ): """ Generates a video based on the given prompt and saves it to the specified path. Parameters: - prompt (str): The description of the video to be generated. - lora_path (str): The path of the LoRA weights to be used. - lora_rank (int): The rank of the LoRA weights. - output_path (str): The path where the generated video will be saved. - num_inference_steps (int): Number of steps for the inference process. More steps can result in better quality. - num_frames (int): Number of frames to generate. CogVideoX1.0 generates 49 frames for 6 seconds at 8 fps, while CogVideoX1.5 produces either 81 or 161 frames, corresponding to 5 seconds or 10 seconds at 16 fps. - width (int): The width of the generated video, applicable only for CogVideoX1.5-5B-I2V - height (int): The height of the generated video, applicable only for CogVideoX1.5-5B-I2V - 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. - dtype (torch.dtype): The data type for computation (default is torch.bfloat16). - generate_type (str): The type of video generation (e.g., 't2v', 'i2v', 'v2v').ยท - seed (int): The seed for reproducibility. - fps (int): The frames per second for the generated video. """ # Distributed local_rank = init_dist(launcher=launcher, port=port) global_rank = dist.get_rank() num_processes = dist.get_world_size() is_main_process = global_rank == 0 torch.manual_seed(seed) random.seed(seed) if is_main_process: os.makedirs(os.path.join(output_path, 'generated_videos'), exist_ok=True) os.makedirs(os.path.join(output_path, 'generated_flow_videos'), exist_ok=True) os.makedirs(os.path.join(output_path, 'flow_warped_videos'), exist_ok=True) # 1. Load the pre-trained CogVideoX pipeline with the specified precision (bfloat16). # add device_map="balanced" in the from_pretrained function and remove the enable_model_cpu_offload() # function to use Multi GPUs. image = None video = None model_name = "cogvideox-5b-i2v".lower() desired_resolution = RESOLUTION_MAP[model_name] if width is None or height is None: height, width = desired_resolution logging.info(f"\033[1mUsing default resolution {desired_resolution} for {model_name}\033[0m") elif (height, width) != desired_resolution: if generate_type == "i2v": # For i2v models, use user-defined width and height logging.warning( f"\033[1;31mThe width({width}) and height({height}) are not recommended for {model_name}. The best resolution is {desired_resolution}.\033[0m" ) """ # Prepare Dataset Class.. """ # image = load_image(image=image_or_video_path) # prompt # first image # camera parameters dataset = I2VFlowDataset_Inference( max_num_frames=num_frames, height=height, width=width, data_root=data_root, max_num_videos=max_num_videos, ) distributed_sampler = DistributedSampler( dataset, num_replicas=num_processes, rank=global_rank, shuffle=False, seed=seed, ) # DataLoaders creation: dataloader = torch.utils.data.DataLoader( dataset, batch_size=1, shuffle=False, sampler=distributed_sampler, num_workers=4, pin_memory=True, drop_last=False, ) """ # Prepare Pipeline """ transformer_lora_config = LoraConfig( r=128, lora_alpha=64, init_lora_weights=True, target_modules=["to_q", "to_k", "to_v", "to_out.0", "norm1.linear", "norm2.linear", "ff.net.2"], ) print(f'Constructing pipeline') pipe_omsm = load_cogvideox_flovd_OMSM_lora_pipeline(omsm_path, backbone_path="THUDM/CogVideoX-5b-I2V", transformer_lora_config=transformer_lora_config, device=local_rank, dtype=dtype) pipe_fvsm = load_cogvideox_flovd_FVSM_controlnet_pipeline(fvsm_path, backbone_path="THUDM/CogVideoX-5b-I2V", device=local_rank, dtype=dtype) print(f'Done loading pipeline') assert pose_type in ['re10k', 'manual'], "Choose other pose_type between ['re10k', 'manual']" if pose_type == 're10k': root_path = "./manual_poses_re10k" else: root_path = "./manual_poses" CameraSampler = SampleManualCam(pose_type=pose_type, root_path=root_path) camera_flow_generator_target = 'finetune.modules.depth_warping.depth_warping.DepthWarping_wrapper' camera_flow_generator = initialize_flow_generator(camera_flow_generator_target).to(local_rank) #-------------------------------------------------------------------------------------------------------- # 2. Set Scheduler. # Can be changed to `CogVideoXDPMScheduler` or `CogVideoXDDIMScheduler`. # We recommend using `CogVideoXDDIMScheduler` for CogVideoX-2B. # using `CogVideoXDPMScheduler` for CogVideoX-5B / CogVideoX-5B-I2V. # pipe.scheduler = CogVideoXDDIMScheduler.from_config(pipe.scheduler.config, timestep_spacing="trailing") pipe_fvsm.scheduler = CogVideoXDPMScheduler.from_config(pipe_fvsm.scheduler.config, timestep_spacing="trailing") pipe_omsm.scheduler = CogVideoXDPMScheduler.from_config(pipe_omsm.scheduler.config, timestep_spacing="trailing") # 3. Enable CPU offload for the model. # turn off if you have multiple GPUs or enough GPU memory(such as H100) and it will cost less time in inference # and enable to("cuda") # pipe.to("cuda") # pipe_fvsm.enable_model_cpu_offload() # pipe_omsm.enable_model_cpu_offload() # pipe_fvsm.enable_sequential_cpu_offload() # pipe_omsm.enable_sequential_cpu_offload() pipe_fvsm.vae.enable_slicing() pipe_fvsm.vae.enable_tiling() pipe_omsm.vae.enable_slicing() pipe_omsm.vae.enable_tiling() dataloader.sampler.set_epoch(1) dist.barrier() output_video_path = os.path.join(output_path, 'generated_videos') output_flow_path = os.path.join(output_path, 'generated_flow_videos') output_warped_video_path = os.path.join(output_path, 'flow_warped_videos') data_iter = iter(dataloader) for step in tqdm(range(0, len(dataloader))): batch = next(data_iter) prompt = batch["prompt"][0] image = batch["image"].to(local_rank) prompt_embedding = batch["prompt_embedding"].to(local_rank) prompt_short = prompt[:20].strip() # if step < 10: # step += 1 # continue # Get Camera flow camparam, cam_name = CameraSampler.sample() # W2C image_torch_255 = ((image.detach().clone()+1)/2. * 255.).squeeze(0) camera_flow_generator_input = get_camera_flow_generator_input(image_torch_255, camparam, device=local_rank, speed=speed) image_torch = ((image_torch_255.unsqueeze(0) / 255.) * 2. - 1.).to(local_rank) with torch.no_grad(): with torch.cuda.amp.autocast(enabled=True, dtype=dtype): # camera_flow, log_dict = camera_flow_generator(image_torch, camera_flow_generator_input) # camera_flow = camera_flow.to(local_rank, dtype) # camera_flow_latent = rearrange(encode_flow(camera_flow, pipe_omsm.vae, flow_scale_factor=[60, 36]), 'b c f h w -> b f c h w').to(local_rank, dtype) flow_latent = pipe_omsm( num_frames=num_frames, height=height, width=width, prompt=None, prompt_embeds=prompt_embedding, image=image, generator=torch.Generator().manual_seed(seed), num_inference_steps=num_inference_steps, use_dynamic_cfg=use_dynamic_cfg, output_type='latent' ).frames[0] object_flow = decode_flow(flow_latent.detach().clone().unsqueeze(0).to(local_rank), pipe_omsm.vae, flow_scale_factor=[60, 36]) # BF,C,H,W if use_flow_integration: # Integrate camera (from 3D warping) and object (from OMSM) flow maps # Using segmentation model will be implemented later.. camera_flow, log_dict = camera_flow_generator(image_torch, camera_flow_generator_input) camera_flow = camera_flow.to(local_rank, dtype) integrated_flow = integrate_flow(camera_flow, object_flow, log_dict['depth_ctxt'], camera_flow_generator, camera_flow_generator_input) integrated_flow_latent = rearrange(encode_flow(integrated_flow, pipe_omsm.vae, flow_scale_factor=[60, 36]), 'b c f h w -> b f c h w').to(local_rank, dtype) else: integrated_flow_latent = rearrange(flow_latent, '(b f) c h w -> b f c h w', b=image.size(0)) # 4. Generate the video frames based on the prompt. # `num_frames` is the Number of frames to generate. video_generate = pipe_fvsm( num_frames=num_frames, height=height, width=width, prompt=None, prompt_embeds=prompt_embedding, image=image, flow_latent=integrated_flow_latent, valid_mask=None, generator=torch.Generator().manual_seed(seed), num_inference_steps=num_inference_steps, controlnet_guidance_start = 0.0, controlnet_guidance_end = controlnet_guidance_end, use_dynamic_cfg=use_dynamic_cfg, ).frames[0] # Save logs # 1) Synthesized flow (object_flow) save_path = os.path.join(output_flow_path, f"{prompt_short}_DCFG-{use_dynamic_cfg}_ContGuide-{controlnet_guidance_end}_{cam_name}.mp4") save_flow(object_flow, filename=save_path, fps=fps) # 2) Flow-Warped Video save_path = os.path.join(output_warped_video_path, f"{prompt_short}_DCFG-{use_dynamic_cfg}_ContGuide-{controlnet_guidance_end}_{cam_name}.mp4") save_flow_warped_video(image_torch_255, object_flow, filename=save_path, fps=fps) # 3) Flow-Cond. Synthesized Video save_path = os.path.join(output_video_path, f"{prompt_short}_DCFG-{use_dynamic_cfg}_ContGuide-{controlnet_guidance_end}_{cam_name}.mp4") export_to_video(video_generate, save_path, fps=fps) dist.barrier() step += 1 #-------------------------------------------------------------------------------------------------- def encode_video(video: torch.Tensor, vae) -> torch.Tensor: # shape of input video: [B, C, F, H, W] video = video.to(vae.device, dtype=vae.dtype) latent_dist = vae.encode(video).latent_dist latent = latent_dist.sample() * vae.config.scaling_factor return latent def encode_flow(flow, vae, flow_scale_factor): # flow: BF,C,H,W # flow_scale_factor [sf_x, sf_y] assert flow.ndim == 4 num_frames, _, height, width = flow.shape # Normalize optical flow # ndim: 4 -> 5 flow = rearrange(flow, '(b f) c h w -> b f c h w', b=1) flow_norm = adaptive_normalize(flow, flow_scale_factor[0], flow_scale_factor[1]) # ndim: 5 -> 4 flow_norm = rearrange(flow_norm, 'b f c h w -> (b f) c h w', b=1) # Duplicate mean value for third channel num_frames, _, H, W = flow_norm.shape flow_norm_extended = torch.empty((num_frames, 3, height, width)).to(flow_norm) flow_norm_extended[:,:2] = flow_norm flow_norm_extended[:,-1:] = flow_norm.mean(dim=1, keepdim=True) flow_norm_extended = rearrange(flow_norm_extended, '(b f) c h w -> b c f h w', f=num_frames) return encode_video(flow_norm_extended, vae) def decode_flow(flow_latent, vae, flow_scale_factor): flow_latent = flow_latent.permute(0, 2, 1, 3, 4) # [batch_size, num_channels, num_frames, height, width] flow_latent = 1 / vae.config.scaling_factor * flow_latent flow = vae.decode(flow_latent).sample # BCFHW # discard third channel (which is a mean value of f_x and f_y) flow = flow[:,:2].detach().clone() # Unnormalize optical flow flow = rearrange(flow, 'b c f h w -> b f c h w') flow = adaptive_unnormalize(flow, flow_scale_factor[0], flow_scale_factor[1]) flow = rearrange(flow, 'b f c h w -> (b f) c h w') return flow # BF,C,H,W def adaptive_normalize(flow, sf_x, sf_y): # x: BFCHW, optical flow assert flow.ndim == 5, 'Set the shape of the flow input as (B, F, C, H, W)' assert sf_x is not None and sf_y is not None b, f, c, h, w = flow.shape max_clip_x = math.sqrt(w/sf_x) * 1.0 max_clip_y = math.sqrt(h/sf_y) * 1.0 flow_norm = flow.detach().clone() flow_x = flow[:, :, 0].detach().clone() flow_y = flow[:, :, 1].detach().clone() flow_x_norm = torch.sign(flow_x) * torch.sqrt(torch.abs(flow_x)/sf_x + 1e-7) flow_y_norm = torch.sign(flow_y) * torch.sqrt(torch.abs(flow_y)/sf_y + 1e-7) flow_norm[:, :, 0] = torch.clamp(flow_x_norm, min=-max_clip_x, max=max_clip_x) flow_norm[:, :, 1] = torch.clamp(flow_y_norm, min=-max_clip_y, max=max_clip_y) return flow_norm def adaptive_unnormalize(flow, sf_x, sf_y): # x: BFCHW, optical flow assert flow.ndim == 5, 'Set the shape of the flow input as (B, F, C, H, W)' assert sf_x is not None and sf_y is not None flow_orig = flow.detach().clone() flow_x = flow[:, :, 0].detach().clone() flow_y = flow[:, :, 1].detach().clone() flow_orig[:, :, 0] = torch.sign(flow_x) * sf_x * (flow_x**2 - 1e-7) flow_orig[:, :, 1] = torch.sign(flow_y) * sf_y * (flow_y**2 - 1e-7) return flow_orig #-------------------------------------------------------------------------------------------------- if __name__ == "__main__": parser = argparse.ArgumentParser(description="Generate a video from a text prompt using CogVideoX") # parser.add_argument("--prompt", type=str, required=True, help="The description of the video to be generated") parser.add_argument("--image_path", type=str, default=None, help="The path of the image to be used as the background of the video",) parser.add_argument("--data_root", type=str, required=True, help="The path of the dataset root",) parser.add_argument("--fvsm_path", type=str, required=True, help="Path of the pre-trained model use") parser.add_argument("--omsm_path", type=str, required=True, help="Path of the pre-trained model use") parser.add_argument("--output_path", type=str, default="./output.mp4", help="The path save generated video") parser.add_argument("--guidance_scale", type=float, default=6.0, help="The scale for classifier-free guidance") parser.add_argument("--num_inference_steps", type=int, default=50, help="Inference steps") parser.add_argument("--num_frames", type=int, default=49, help="Number of steps for the inference process") parser.add_argument("--width", type=int, default=None, help="The width of the generated video") parser.add_argument("--height", type=int, default=None, help="The height of the generated video") parser.add_argument("--fps", type=int, default=16, help="The frames per second for the generated video") parser.add_argument("--num_videos_per_prompt", type=int, default=1, help="Number of videos to generate per prompt") parser.add_argument("--dtype", type=str, default="bfloat16", help="The data type for computation") parser.add_argument("--seed", type=int, default=42, help="The seed for reproducibility") parser.add_argument("--controlnet_guidance_end", type=float, default=0.4, help="Controlnet guidance end during sampling") parser.add_argument("--max_num_videos", type=int, default=None, help="# of videos for inference") parser.add_argument("--use_dynamic_cfg", action='store_true') parser.add_argument("--pose_type", type=str, default='manual', help="pose type in the inference time") parser.add_argument("--speed", type=float, default=0.5, help="pose type in the inference time") parser.add_argument("--use_flow_integration", action='store_true') # DDP args parser.add_argument("--launcher", type=str, choices=["pytorch", "slurm"], default="pytorch") parser.add_argument("--world_size", default=1, type=int, help="number of the distributed processes.") parser.add_argument('--local-rank', type=int, default=-1, help='Replica rank on the current node. This field is required ' 'by `torch.distributed.launch`.') parser.add_argument("--global_seed", default=42, type=int, help="seed") parser.add_argument("--port", type=int) parser.add_argument("--local_rank", type=int, help="Local rank. Necessary for using the torch.distributed.launch utility.") args = parser.parse_args() dtype = torch.float16 if args.dtype == "float16" else torch.bfloat16 generate_video( # prompt=args.prompt, launcher=args.launcher, port=args.port, data_root=args.data_root, fvsm_path=args.fvsm_path, omsm_path=args.omsm_path, output_path=args.output_path, num_frames=args.num_frames, width=args.width, height=args.height, image_path=args.image_path, num_inference_steps=args.num_inference_steps, guidance_scale=args.guidance_scale, num_videos_per_prompt=args.num_videos_per_prompt, dtype=dtype, seed=args.seed, fps=args.fps, controlnet_guidance_end=args.controlnet_guidance_end, max_num_videos=args.max_num_videos, use_dynamic_cfg=args.use_dynamic_cfg, pose_type=args.pose_type, speed=args.speed, use_flow_integration=args.use_flow_integration, )