import copy import gc import logging import math import os import random import sys from contextlib import contextmanager import numpy as np import torch import torch.cuda.amp as amp import torchvision.transforms.functional as TF from tqdm import tqdm from .modules.clip import CLIPModel from .modules.t5 import T5EncoderModel from .modules.vae import WanVAE from .modules.model import WanModel from .utils.fm_solvers import (FlowDPMSolverMultistepScheduler, get_sampling_sigmas, retrieve_timesteps) from .utils.fm_solvers_unipc import FlowUniPCMultistepScheduler from .utils.audio_utils import preprocess_audio, resample_audio from .utils.infer_utils import create_null_audio_ref_features, \ expand_face_mask_flexible, gen_inference_masks, expand_bbox_and_crop_image, count_parameters, \ gen_smooth_transition_mask_for_dit, process_audio_features, process_audio_features class WanAF2V: def __init__( self, config, checkpoint_dir, device_id=0, rank=0, t5_fsdp=False, dit_fsdp=False, t5_cpu=False, init_on_cpu=True, use_gradient_checkpointing=False, post_trained_checkpoint_path=None, dit_config=None, crop_image_size=224, use_half=False, ): r""" Initializes the image-to-video generation model components. Args: config (EasyDict): Object containing model parameters initialized from config.py checkpoint_dir (`str`): Path to directory containing model checkpoints device_id (`int`, *optional*, defaults to 0): Id of target GPU device rank (`int`, *optional*, defaults to 0): Process rank for distributed training t5_fsdp (`bool`, *optional*, defaults to False): Enable FSDP sharding for T5 model dit_fsdp (`bool`, *optional*, defaults to False): Enable FSDP sharding for DiT model t5_cpu (`bool`, *optional*, defaults to False): Whether to place T5 model on CPU. Only works without t5_fsdp. init_on_cpu (`bool`, *optional*, defaults to True): Enable initializing Transformer Model on CPU. Only works without FSDP or USP. post_trained_checkpoint_path (`str`, *optional*, defaults to None): Path to the post-trained checkpoint file. If provided, model will be loaded from this checkpoint. use_half (`bool`, *optional*, defaults to False): Whether to use half precision (float16/bfloat16) for model inference. Reduces memory usage. """ # 如果 device_id 为 -1 或 CUDA 不可用，使用 CPU if device_id == -1 or not torch.cuda.is_available(): self.device = torch.device("cpu") else: self.device = torch.device(f"cuda:{device_id}") self.config = config self.rank = rank self.t5_cpu = t5_cpu self.use_half = use_half self.num_train_timesteps = config.num_train_timesteps self.param_dtype = config.param_dtype # If half precision is enabled, force use float16 for faster speed if use_half: self.half_dtype = torch.float16 logging.info(f"Half precision enabled, using dtype: {self.half_dtype} (forced float16 for faster inference)") else: self.half_dtype = torch.float32 self.text_encoder = T5EncoderModel( text_len=config.text_len, dtype=config.t5_dtype, device=torch.device('cpu'), checkpoint_path=os.path.join(checkpoint_dir, config.t5_checkpoint), tokenizer_path=os.path.join(checkpoint_dir, config.t5_tokenizer), shard_fn=None, ) # Output model parameter count model_param_count = count_parameters(self.text_encoder.model) logging.info(f"Text Model parameters: {model_param_count}M") self.crop_image_size = crop_image_size self.vae_stride = config.vae_stride self.patch_size = config.patch_size self.vae = WanVAE( vae_pth=os.path.join(checkpoint_dir, config.vae_checkpoint), device=self.device) # Output model parameter count model_param_count = count_parameters(self.vae.model) logging.info(f"VAE Model parameters: {model_param_count}M") self.clip = CLIPModel( dtype=config.clip_dtype, device=self.device, checkpoint_path=os.path.join(checkpoint_dir, config.clip_checkpoint), tokenizer_path=os.path.join(checkpoint_dir, config.clip_tokenizer)) # Output model parameter count model_param_count = count_parameters(self.clip.model) logging.info(f"CLIP Model parameters: {model_param_count}M") logging.info(f"Creating WanModel from {checkpoint_dir}") if post_trained_checkpoint_path: try: if rank == 0: print(f"Loading post-trained model from {post_trained_checkpoint_path}") # Load model config from original directory # TODO: config can also be specified in current directory # config_dict = json.load(open(os.path.join(checkpoint_dir, 'config.json'))) config_dict = dit_config self.model = WanModel.from_config(config_dict) # All cards directly load safetensors #model_state = load_file(post_trained_checkpoint_path) checkpoint = torch.load(post_trained_checkpoint_path, map_location='cpu', weights_only=True) model_state = checkpoint['model'] self.model.load_state_dict(model_state) if rank == 0: print(f"safertensors have been loaded: {post_trained_checkpoint_path}") except Exception as e: if rank == 0: print(f"Error loading post-trained model: {e}") raise e else: self.model = WanModel.from_pretrained(checkpoint_dir) self.model.eval().requires_grad_(False) # Output model parameter count model_param_count = count_parameters(self.model) logging.info(f"DiT Model parameters: {model_param_count}M") # Enable gradient checkpointing (only effective during training) if use_gradient_checkpointing: self.model.enable_gradient_checkpointing() # Assume WanModel supports this method logging.info("Gradient checkpointing enabled for WanModel") self.sp_size = 1 if not init_on_cpu: self.model.to(self.device) # If half precision is enabled, convert model to half precision if use_half: try: self.model = self.model.to(dtype=self.half_dtype) logging.info(f"Model converted to {self.half_dtype} precision") except Exception as e: logging.warning(f"Failed to convert model to half precision: {e}. Continuing with float32.") self.use_half = False self.half_dtype = torch.float32 self.sample_neg_prompt = config.sample_neg_prompt def generate( self, input_prompt, img, audio=None, # New audio input max_area=720 * 1280, frame_num=81, shift=5.0, sample_solver='unipc', sampling_steps=40, guide_scale=5.0, n_prompt="", seed=-1, offload_model=True, cfg_zero=False, zero_init_steps=0, face_processor=None, # InsightFace model img_path=None, # For InsightFace use audio_paths=None, # New: audio path list, supports multiple audio files task_key=None, mode="pad", # Audio processing mode: "pad" or "concat" trim_to_4s=False, # Fast mode: trim audio to 4 seconds ): r""" Generates video frames from input image and text prompt using diffusion process. Args: input_prompt (`str`): Text prompt for content generation. img (PIL.Image.Image): Input image tensor. Shape: [3, H, W] max_area (`int`, *optional*, defaults to 720*1280): Maximum pixel area for latent space calculation. Controls video resolution scaling frame_num (`int`, *optional*, defaults to 81): How many frames to sample from a video. The number should be 4n+1 shift (`float`, *optional*, defaults to 5.0): Noise schedule shift parameter. Affects temporal dynamics [NOTE]: If you want to generate a 480p video, it is recommended to set the shift value to 3.0. sample_solver (`str`, *optional*, defaults to 'unipc'): Solver used to sample the video. sampling_steps (`int`, *optional*, defaults to 40): Number of diffusion sampling steps. Higher values improve quality but slow generation guide_scale (`float`, *optional*, defaults 5.0): Classifier-free guidance scale. Controls prompt adherence vs. creativity n_prompt (`str`, *optional*, defaults to ""): Negative prompt for content exclusion. If not given, use `config.sample_neg_prompt` seed (`int`, *optional*, defaults to -1): Random seed for noise generation. If -1, use random seed offload_model (`bool`, *optional*, defaults to True): If True, offloads models to CPU during generation to save VRAM cfg_zero (`bool`, *optional*, defaults to False): Whether to use adaptive CFG-Zero guidance instead of fixed guidance scale zero_init_steps (`int`, *optional*, defaults to 0): Number of initial steps to use zero guidance when using cfg_zero Returns: torch.Tensor: Generated video frames tensor. Dimensions: (C, N H, W) where: - C: Color channels (3 for RGB) - N: Number of frames (81) - H: Frame height (from max_area) - W: Frame width from max_area) """ img = TF.to_tensor(img).sub_(0.5).div_(0.5).to(self.device) # If half precision is enabled, convert input image to half precision if self.use_half: img = img.to(dtype=self.half_dtype) # Save audio_paths parameter for later use self.audio_paths = audio_paths # Use passed frame_num directly (it's already calculated outside in concat mode) F = frame_num print(f"Using frame number: {F} (mode: {mode})") h, w = img.shape[1:] print(f"Input image size: {h}, {w}, {max_area}") aspect_ratio = h / w lat_h = round( np.sqrt(max_area * aspect_ratio) // self.vae_stride[1] // self.patch_size[1] * self.patch_size[1]) lat_w = round( np.sqrt(max_area / aspect_ratio) // self.vae_stride[2] // self.patch_size[2] * self.patch_size[2]) h = lat_h * self.vae_stride[1] w = lat_w * self.vae_stride[2] # Save original aspect ratio original_aspect_ratio = h / w # Find nearest size combination divisible by 32 # First try adjusting based on width w_candidate1 = (w // 32) * 32 if w_candidate1 == 0: w_candidate1 = 32 h_candidate1 = (int(w_candidate1 * original_aspect_ratio) // 32) * 32 if h_candidate1 == 0: h_candidate1 = 32 # Then try adjusting based on height h_candidate2 = (h // 32) * 32 if h_candidate2 == 0: h_candidate2 = 32 w_candidate2 = (int(h_candidate2 / original_aspect_ratio) // 32) * 32 if w_candidate2 == 0: w_candidate2 = 32 # Select combination with smallest difference from original size diff1 = abs(w_candidate1 - w) + abs(h_candidate1 - h) diff2 = abs(w_candidate2 - w) + abs(h_candidate2 - h) if diff1 <= diff2: w, h = w_candidate1, h_candidate1 else: w, h = w_candidate2, h_candidate2 # Recalculate lat_h and lat_w to maintain consistency lat_h = h // self.vae_stride[1] lat_w = w // self.vae_stride[2] print(f"Processed image size: {h}, {w}, {lat_h}, {lat_w}") # Create noise after adjusting lat_h and lat_w latent_frame_num = (F - 1) // self.vae_stride[0] + 1 seed = seed if seed >= 0 else random.randint(0, sys.maxsize) seed_g = torch.Generator(device=self.device) seed_g.manual_seed(seed) # Select dtype based on whether half precision is enabled noise_dtype = self.half_dtype if self.use_half else torch.float32 noise = torch.randn(16, latent_frame_num, lat_h, lat_w, dtype=noise_dtype, generator=seed_g, device=self.device) print(f"noise shape: {noise.shape}, dtype: {noise.dtype}") max_seq_len = ((F - 1) // self.vae_stride[0] + 1) * lat_h * lat_w // ( self.patch_size[1] * self.patch_size[2]) max_seq_len = int(math.ceil(max_seq_len / self.sp_size)) * self.sp_size print(f"Max seq_len: {max_seq_len}") if n_prompt == "": n_prompt = self.sample_neg_prompt # preprocess if not self.t5_cpu: self.text_encoder.model.to(self.device) context = self.text_encoder([input_prompt], self.device) context_null = self.text_encoder([n_prompt], self.device) if offload_model: self.text_encoder.model.cpu() else: context = self.text_encoder([input_prompt], torch.device('cpu')) context_null = self.text_encoder([n_prompt], torch.device('cpu')) context = [t.to(self.device) for t in context] context_null = [t.to(self.device) for t in context_null] self.clip.model.to(self.device) with torch.no_grad(): # Ensure input image uses correct precision img_input = img[:, None, :, :] if self.use_half: img_input = img_input.to(dtype=self.half_dtype) clip_context = self.clip.visual([img_input]) # If half precision is enabled, convert clip_context to half precision if self.use_half: clip_context = clip_context.to(dtype=self.half_dtype) """ Start of i2v mask and ref latent construction logic """ # Select mask dtype based on whether half precision is enabled mask_dtype = self.half_dtype if self.use_half else torch.float32 msk = torch.ones(1, F, lat_h, lat_w, device=self.device, dtype=mask_dtype) # Use F corresponding to original frame number msk[:, 1:] = 0 msk = torch.concat([ torch.repeat_interleave(msk[:, 0:1], repeats=4, dim=1), msk[:, 1:] ], dim=1) msk = msk.view(1, msk.shape[1] // 4, 4, lat_h, lat_w) msk = msk.transpose(1, 2)[0] # Latent containing first frame ref with torch.no_grad(): # VAE encoding requires float32, but we can convert result to half precision img_for_vae = img[None].cpu() if img_for_vae.dtype != torch.float32: img_for_vae = img_for_vae.float() y = self.vae.encode([ torch.concat([ torch.nn.functional.interpolate( img_for_vae, size=(h, w), mode='bicubic').transpose( 0, 1), torch.zeros(3, F-1, h, w, dtype=torch.float32) ], dim=1).to(self.device) ])[0] # If half precision is enabled, convert VAE encoding result to half precision if self.use_half: y = y.to(dtype=self.half_dtype) print(f"y shape after VAE encode: {y.shape}, dtype: {y.dtype}") # Initialize face_clip_context variable face_clip_context = None if face_processor is not None: w_scale_factor = 1.2 h_scale_factor = 1.1 # Determine number of faces needed based on number of audio files if hasattr(self, 'audio_paths') and self.audio_paths: # Use new audio_paths parameter n_faces_needed = len(self.audio_paths) print(f"number of faces needed: {n_faces_needed}") else: # Only one audio or no audio n_faces_needed = 1 print(f"only one audio or no audio, number of faces needed: {n_faces_needed}") # Use modified infer method, pass n parameter face_info = face_processor.infer(img_path, n=n_faces_needed) # Get actual number of faces detected n = len(face_info['masks']) # Get sample count print(f"number of faces detected: {n}") # Strictly arrange left-to-right, no random sorting # Assume face detection results are already arranged left-to-right print(f"the face order is set to left-to-right: {list(range(n))}") masks = face_info['masks'] print(f"video with face processor, scale up w={w_scale_factor}, h={h_scale_factor} ###") # First expand mask, then perform smooth transition processing expanded_masks = [] for mask in masks: expanded_mask = expand_face_mask_flexible( torch.from_numpy(mask).to(self.device, dtype=img.dtype).clone(), width_scale_factor=w_scale_factor, height_scale_factor=h_scale_factor ) expanded_masks.append(expanded_mask) global_mask = torch.zeros_like(expanded_masks[0]) for mask in expanded_masks: global_mask += mask # Select mask dtype based on whether half precision is enabled mask_dtype = self.half_dtype if self.use_half else torch.float32 dit_mask = gen_smooth_transition_mask_for_dit( global_mask, lat_h, lat_w, F, device=self.device, mask_dtype=mask_dtype, target_translate=(0, 0), # Random left or right target_scale=1 # No longer scaling in gen_smooth_transition_mask_for_dit ) y = torch.cat([dit_mask, y], dim=0) # Resize expanded_masks to VAE encoded latent size resized_masks = [] with torch.no_grad(): for mask in expanded_masks: # Resize mask from original size directly to VAE encoded latent size latent_mask = torch.nn.functional.interpolate( mask.unsqueeze(0).unsqueeze(0), # [1, 1, H, W] size=(lat_h // 2, lat_w // 2), # Directly resize to latent size mode='bilinear', align_corners=False ).squeeze(0).squeeze(0) resized_masks.append(latent_mask) # Calculate number of frames after VAE encoding latent_frame_num = (F - 1) // self.vae_stride[0] + 1 inference_masks = gen_inference_masks( resized_masks, (lat_h // 2, lat_w // 2), # Use VAE encoded latent size num_frames=latent_frame_num # Use number of frames after VAE encoding ) bboxes = face_info['bboxes'] # Loop through all faces, extract CLIP features face_clip_context_list = [] for i, bbox in enumerate(bboxes): try: # Directly process bbox, assume format is [x, y, w, h] bbox_x, bbox_y, bbox_w, bbox_h = bbox bbox_converted = [bbox_x, bbox_y, bbox_x + bbox_w, bbox_y + bbox_h] # Convert to [x1, y1, x2, y2] # Use encapsulated function for bbox expansion and image cropping, but cropping uses original bbox cropped_face, adjusted_bbox = expand_bbox_and_crop_image( img, bbox_converted, width_scale_factor=1, height_scale_factor=1 ) cropped_face = cropped_face.to(self.device) # If half precision is enabled, convert cropped face to half precision if self.use_half: cropped_face = cropped_face.to(dtype=self.half_dtype) with torch.no_grad(): # CLIP encoding method face_clip_context = self.clip.visual([cropped_face[:, None, :, :]]) # If half precision is enabled, convert face_clip_context to half precision if self.use_half: face_clip_context = face_clip_context.to(dtype=self.half_dtype) # Add to list face_clip_context_list.append(face_clip_context) except Exception as e: print(f"error on face {i+1}: {e}") continue print(f"face feature extraction loop completed, successfully processed {len(face_clip_context_list)} faces") # For backward compatibility, keep original variable name face_clip_context = face_clip_context_list[0] if face_clip_context_list else None else: y = torch.concat([msk, y]) # Process multiple audio files using utility function audio_feat_list = process_audio_features( audio_paths=audio_paths, audio=audio, mode=mode, F=F, frame_num=frame_num, task_key=task_key, fps=self.config.fps, wav2vec_model=self.config.wav2vec, vocal_separator_model=self.config.vocal_separator_path, audio_output_dir=self.config.audio_output_dir, device=self.device, use_half=self.use_half, half_dtype=self.half_dtype, preprocess_audio=preprocess_audio, resample_audio=resample_audio, trim_to_4s=trim_to_4s, ) # Prepare audio_ref_features - new list mode audio_ref_features = { "ref_face": face_clip_context, # face clip features (backward compatible) } # Build face feature list and audio feature list ref_face_list = face_clip_context_list.copy() if face_clip_context_list else [] audio_list = audio_feat_list.copy() if audio_feat_list else [] # Add lists to audio_ref_features audio_ref_features["ref_face_list"] = ref_face_list audio_ref_features["audio_list"] = audio_list @contextmanager def noop_no_sync(): yield no_sync = getattr(self.model, 'no_sync', noop_no_sync) if offload_model: self.clip.model.cpu() torch.cuda.empty_cache() # evaluation mode with amp.autocast(dtype=self.param_dtype), torch.no_grad(), no_sync(): if sample_solver == 'unipc': sample_scheduler = FlowUniPCMultistepScheduler( num_train_timesteps=self.num_train_timesteps, shift=5, use_dynamic_shifting=False) sample_scheduler.set_timesteps( sampling_steps, device=self.device, shift=shift) timesteps = sample_scheduler.timesteps elif sample_solver == 'dpm++': sample_scheduler = FlowDPMSolverMultistepScheduler( num_train_timesteps=self.num_train_timesteps, shift=1, use_dynamic_shifting=False) sampling_sigmas = get_sampling_sigmas(sampling_steps, shift) timesteps, _ = retrieve_timesteps( sample_scheduler, device=self.device, sigmas=sampling_sigmas) else: raise NotImplementedError("Unsupported solver.") # sample videos latent = noise # Ensure latent uses correct precision if self.use_half and latent.dtype != self.half_dtype: latent = latent.to(dtype=self.half_dtype) # If half precision is enabled, convert context to half precision if self.use_half: context = [c.to(dtype=self.half_dtype) for c in context] context_null = [c.to(dtype=self.half_dtype) for c in context_null] # for cfg - use complete audio_ref_features arg_c = { 'context': [context[0]], 'clip_fea': clip_context, 'seq_len': max_seq_len, 'y': [y], 'audio_ref_features': audio_ref_features } # Create different types of null condition parameters arg_null = { 'context': context_null, 'clip_fea': clip_context, 'seq_len': max_seq_len, 'y': [y], 'audio_ref_features': create_null_audio_ref_features(audio_ref_features) } # Add inference-generated masks to model parameters if face_processor is not None and 'inference_masks' in locals(): # Use deep copy to ensure each call uses independent mask data arg_c['face_mask_list'] = copy.deepcopy(inference_masks['face_mask_list']) arg_null['face_mask_list'] = copy.deepcopy(inference_masks['face_mask_list']) if offload_model: torch.cuda.empty_cache() self.model.to(self.device) masks_flattened = False # Set to False for first time for i, t in enumerate(tqdm(timesteps)): # Ensure latent is on correct device and precision latent = latent.to(self.device) if self.use_half and latent.dtype != self.half_dtype: latent = latent.to(dtype=self.half_dtype) latent_model_input = [latent] timestep = [t] timestep = torch.stack(timestep).to(self.device) model_output_cond = self.model( latent_model_input, t=timestep, masks_flattened=masks_flattened, **arg_c ) noise_pred_cond = model_output_cond[0] # Ensure noise_pred_cond uses correct precision if self.use_half and noise_pred_cond.dtype != self.half_dtype: noise_pred_cond = noise_pred_cond.to(dtype=self.half_dtype) noise_pred_cond = noise_pred_cond.to( torch.device('cpu') if offload_model else self.device) if offload_model: torch.cuda.empty_cache() if not cfg_zero: model_output_uncond = self.model( latent_model_input, t=timestep, masks_flattened=masks_flattened, **arg_null ) noise_pred_uncond = model_output_uncond[0] # Ensure noise_pred_uncond uses correct precision if self.use_half and noise_pred_uncond.dtype != self.half_dtype: noise_pred_uncond = noise_pred_uncond.to(dtype=self.half_dtype) noise_pred_uncond = noise_pred_uncond.to( torch.device('cpu') if offload_model else self.device) else: noise_pred_uncond = None masks_flattened = True # No need to execute later if offload_model: torch.cuda.empty_cache() if cfg_zero: noise_pred = noise_pred_cond else: noise_pred = noise_pred_uncond + guide_scale * (noise_pred_cond - noise_pred_uncond) # Ensure noise_pred uses correct precision if self.use_half and noise_pred.dtype != self.half_dtype: noise_pred = noise_pred.to(dtype=self.half_dtype) latent = latent.to( torch.device('cpu') if offload_model else self.device ) # scheduler.step may require float32, so temporarily convert noise_pred_for_scheduler = noise_pred.float() if self.use_half else noise_pred latent_for_scheduler = latent.float() if self.use_half else latent temp_x0 = sample_scheduler.step( noise_pred_for_scheduler.unsqueeze(0), t, latent_for_scheduler.unsqueeze(0), return_dict=False, generator=seed_g)[0] # Convert back to half precision (if enabled) if self.use_half: temp_x0 = temp_x0.to(dtype=self.half_dtype) latent = temp_x0.squeeze(0) x0 = [latent.to(self.device)] del latent_model_input, timestep if offload_model: self.model.cpu() torch.cuda.empty_cache() # VAE decoding requires float32, so temporarily convert x0_for_decode = [x.float() if self.use_half else x for x in x0] videos = self.vae.decode(x0_for_decode) result_videos = videos[0] del noise, latent del sample_scheduler if offload_model: gc.collect() torch.cuda.synchronize() return result_videos if self.rank == 0 else None