DFAGWE/infinitetalk2
0
1# Copyright 2024-2025 The Alibaba Wan Team Authors. All rights reserved.2import gc3from inspect import ArgSpec4import logging5import json6import math7import importlib8import os9import random10import sys11import types12from contextlib import contextmanager13from functools import partial14from PIL import Image15 16import numpy as np17import torch18import torch.cuda.amp as amp19import torch.distributed as dist20import torchvision.transforms as transforms21import torch.nn.functional as F22import torch.nn as nn23from tqdm import tqdm24from diffusers.models.modeling_utils import no_init_weights, ContextManagers25import accelerate26 27from .distributed.fsdp import shard_model28from .modules.clip import CLIPModel29from .modules.multitalk_model import WanModel, WanLayerNorm, WanRMSNorm30from .modules.t5 import T5EncoderModel, T5LayerNorm, T5RelativeEmbedding31from .modules.vae import WanVAE, CausalConv3d, RMS_norm, Upsample32from .utils.multitalk_utils import MomentumBuffer, adaptive_projected_guidance, match_and_blend_colors33from src.vram_management import AutoWrappedQLinear, AutoWrappedLinear, AutoWrappedModule, enable_vram_management34from wan.utils.utils import convert_video_to_h264, extract_specific_frames, get_video_codec35from wan.wan_lora import WanLoraWrapper36 37from safetensors.torch import load_file38from optimum.quanto import quantize, freeze, qint8,requantize39import optimum.quanto.nn.qlinear as qlinear40 41def torch_gc():42 torch.cuda.empty_cache()43 torch.cuda.ipc_collect()44 45def to_param_dtype_fp32only(model, param_dtype):46 for module in model.modules():47 for name, param in module.named_parameters(recurse=False):48 if param.dtype == torch.float32 and param.__class__.__name__ not in ['WeightQBytesTensor']:49 param.data = param.data.to(param_dtype)50 for name, buf in module.named_buffers(recurse=False):51 if buf.dtype == torch.float32 and buf.__class__.__name__ not in ['WeightQBytesTensor']:52 module._buffers[name] = buf.to(param_dtype)53 54def resize_and_centercrop(cond_image, target_size):55 """56 Resize image or tensor to the target size without padding.57 """58 59 # Get the original size60 if isinstance(cond_image, torch.Tensor):61 _, orig_h, orig_w = cond_image.shape62 else:63 orig_h, orig_w = cond_image.height, cond_image.width64 65 target_h, target_w = target_size66 67 # Calculate the scaling factor for resizing68 scale_h = target_h / orig_h69 scale_w = target_w / orig_w70 71 # Compute the final size72 scale = max(scale_h, scale_w)73 final_h = math.ceil(scale * orig_h)74 final_w = math.ceil(scale * orig_w)75 76 # Resize77 if isinstance(cond_image, torch.Tensor):78 if len(cond_image.shape) == 3:79 cond_image = cond_image[None]80 resized_tensor = nn.functional.interpolate(cond_image, size=(final_h, final_w), mode='nearest').contiguous() 81 # crop82 cropped_tensor = transforms.functional.center_crop(resized_tensor, target_size) 83 cropped_tensor = cropped_tensor.squeeze(0)84 else:85 resized_image = cond_image.resize((final_w, final_h), resample=Image.BILINEAR)86 resized_image = np.array(resized_image)87 # tensor and crop88 resized_tensor = torch.from_numpy(resized_image)[None, ...].permute(0, 3, 1, 2).contiguous()89 cropped_tensor = transforms.functional.center_crop(resized_tensor, target_size)90 cropped_tensor = cropped_tensor[:, :, None, :, :] 91 92 return cropped_tensor93 94 95def timestep_transform(96 t,97 shift=5.0,98 num_timesteps=1000,99):100 t = t / num_timesteps101 # shift the timestep based on ratio102 new_t = shift * t / (1 + (shift - 1) * t)103 new_t = new_t * num_timesteps104 return new_t105 106 107 108class InfiniteTalkPipeline:109 110 def __init__(111 self,112 config,113 checkpoint_dir,114 quant_dir=None,115 device_id=0,116 rank=0,117 t5_fsdp=False,118 dit_fsdp=False,119 use_usp=False,120 t5_cpu=False,121 init_on_cpu=True,122 num_timesteps=1000,123 use_timestep_transform=True,124 lora_dir=None,125 lora_scales=None,126 quant = None,127 dit_path = None,128 infinitetalk_dir=None,129 ):130 r"""131 Initializes the image-to-video generation model components.132 133 Args:134 config (EasyDict):135 Object containing model parameters initialized from config.py136 checkpoint_dir (`str`):137 Path to directory containing model checkpoints138 device_id (`int`, *optional*, defaults to 0):139 Id of target GPU device140 rank (`int`, *optional*, defaults to 0):141 Process rank for distributed training142 t5_fsdp (`bool`, *optional*, defaults to False):143 Enable FSDP sharding for T5 model144 dit_fsdp (`bool`, *optional*, defaults to False):145 Enable FSDP sharding for DiT model146 use_usp (`bool`, *optional*, defaults to False):147 Enable distribution strategy of USP.148 t5_cpu (`bool`, *optional*, defaults to False):149 Whether to place T5 model on CPU. Only works without t5_fsdp.150 init_on_cpu (`bool`, *optional*, defaults to True):151 Enable initializing Transformer Model on CPU. Only works without FSDP or USP.152 quant (`str`, *optional*, defaults to None):153 Quantization type, must be 'int8' or 'fp8'.154 """155 if quant is not None and quant not in ("int8", "fp8"):156 raise ValueError("quant must be 'int8', 'fp8', or None(default fp32 model)")157 self.device = torch.device(f"cuda:{device_id}")158 self.config = config159 self.rank = rank160 self.use_usp = use_usp161 self.t5_cpu = t5_cpu162 163 self.num_train_timesteps = config.num_train_timesteps164 self.param_dtype = config.param_dtype165 166 shard_fn = partial(shard_model, device_id=device_id)167 168 self.text_encoder = T5EncoderModel(169 text_len=config.text_len,170 dtype=config.t5_dtype,171 device=torch.device('cpu'),172 checkpoint_path=os.path.join(checkpoint_dir, config.t5_checkpoint),173 tokenizer_path=os.path.join(checkpoint_dir, config.t5_tokenizer),174 shard_fn=shard_fn if t5_fsdp else None,175 quant=quant,176 quant_dir=os.path.dirname(quant_dir) if quant_dir is not None else None,177 )178 179 self.vae_stride = config.vae_stride180 self.patch_size = config.patch_size181 self.vae = WanVAE(182 vae_pth=os.path.join(checkpoint_dir, config.vae_checkpoint),183 device=self.device)184 185 self.clip = CLIPModel(186 dtype=config.clip_dtype,187 device=self.device,188 checkpoint_path=os.path.join(checkpoint_dir,189 config.clip_checkpoint),190 tokenizer_path=os.path.join(checkpoint_dir, config.clip_tokenizer))191 192 logging.info(f"Creating WanModel from {checkpoint_dir}")193 194 if quant is not None:195 logging.info(f"Loading Quantized MultiTalk from {quant_dir}")196 with torch.device('meta'):197 wan_config = json.load(open(os.path.join(checkpoint_dir, "config.json")))198 self.model = WanModel(weight_init=False,**wan_config)199 torch_gc()200 model_state_dict = load_file(quant_dir)201 map_json_path = os.path.join(quant_dir.replace('safetensors', 'json'))202 self.model.init_freqs()203 with open(map_json_path, "r") as f:204 quantization_map = json.load(f)205 requantize(self.model, model_state_dict, quantization_map, device='cpu')206 else:207 if dit_path is None:208 init_contexts = [no_init_weights()]209 init_contexts.append(accelerate.init_empty_weights())210 wan_config = json.load(open(os.path.join(checkpoint_dir, "config.json")))211 self.model = WanModel(weight_init=False,**wan_config).to(dtype=self.param_dtype)212 weight_files = [f"{checkpoint_dir}/diffusion_pytorch_model-00001-of-00007.safetensors", 213 f"{checkpoint_dir}/diffusion_pytorch_model-00002-of-00007.safetensors", 214 f"{checkpoint_dir}/diffusion_pytorch_model-00003-of-00007.safetensors", 215 f"{checkpoint_dir}/diffusion_pytorch_model-00004-of-00007.safetensors",216 f"{checkpoint_dir}/diffusion_pytorch_model-00005-of-00007.safetensors", 217 f"{checkpoint_dir}/diffusion_pytorch_model-00006-of-00007.safetensors", 218 f"{checkpoint_dir}/diffusion_pytorch_model-00007-of-00007.safetensors",219 f"{infinitetalk_dir}"]220 merged_state_dict = {}221 for weight_file in weight_files:222 sd = load_file(weight_file)223 merged_state_dict.update(sd)224 self.model.load_state_dict(merged_state_dict)225 226 else:227 init_contexts = [no_init_weights()]228 init_contexts.append(accelerate.init_empty_weights())229 with ContextManagers(init_contexts):230 wan_config = json.load(open(os.path.join(checkpoint_dir, "config.json")))231 self.model = WanModel(weight_init=False,**wan_config)232 checkpoint_weights = torch.load(dit_path, map_location='cpu')233 self.model.load_state_dict(checkpoint_weights['state_dict'])234 logging.info(f"loading infinitetalk weights {checkpoint_dir}")235 236 self.model.eval().requires_grad_(False)237 238 to_param_dtype_fp32only(self.model, self.param_dtype)239 if lora_dir is not None and quant is None :240 lora_wrapper = WanLoraWrapper(self.model)241 for lora_path, lora_scale in zip(lora_dir, lora_scales):242 lora_name = lora_wrapper.load_lora(lora_path)243 lora_wrapper.apply_lora(lora_name, lora_scale, param_dtype=self.param_dtype, device=self.device)244 245 246 247 248 if t5_fsdp or dit_fsdp or use_usp:249 init_on_cpu = False250 if use_usp:251 from xfuser.core.distributed import get_sequence_parallel_world_size252 253 from .distributed.xdit_context_parallel import (254 usp_dit_forward_multitalk,255 usp_attn_forward_multitalk,256 usp_crossattn_multi_forward_multitalk257 )258 for block in self.model.blocks:259 block.self_attn.forward = types.MethodType(260 usp_attn_forward_multitalk, block.self_attn)261 block.audio_cross_attn.forward = types.MethodType(262 usp_crossattn_multi_forward_multitalk, block.audio_cross_attn)263 self.model.forward = types.MethodType(usp_dit_forward_multitalk, self.model)264 self.sp_size = get_sequence_parallel_world_size()265 else:266 self.sp_size = 1267 268 269 270 if dist.is_initialized():271 dist.barrier()272 if dit_fsdp:273 self.model = shard_fn(self.model)274 else:275 if not init_on_cpu:276 self.model.to(self.device)277 278 self.sample_neg_prompt = config.sample_neg_prompt279 self.num_timesteps = num_timesteps280 self.use_timestep_transform = use_timestep_transform281 282 self.cpu_offload = False283 self.model_names = ["model"]284 self.vram_management = False285 286 def add_noise(287 self,288 original_samples: torch.FloatTensor,289 noise: torch.FloatTensor,290 timesteps: torch.IntTensor,291 ) -> torch.FloatTensor:292 """293 compatible with diffusers add_noise()294 """295 timesteps = timesteps.float() / self.num_timesteps296 timesteps = timesteps.view(timesteps.shape + (1,) * (len(noise.shape)-1))297 298 return (1 - timesteps) * original_samples + timesteps * noise299 300 def enable_vram_management(self, num_persistent_param_in_dit=None):301 dtype = next(iter(self.model.parameters())).dtype302 enable_vram_management(303 self.model,304 module_map={305 qlinear.QLinear: AutoWrappedQLinear,306 torch.nn.Linear: AutoWrappedLinear,307 torch.nn.Conv3d: AutoWrappedModule,308 torch.nn.LayerNorm: AutoWrappedModule,309 WanLayerNorm: AutoWrappedModule,310 WanRMSNorm: AutoWrappedModule,311 },312 module_config=dict(313 offload_dtype=dtype,314 offload_device="cpu",315 onload_dtype=dtype,316 onload_device=self.device,317 computation_dtype=self.param_dtype,318 computation_device=self.device,319 ),320 max_num_param=num_persistent_param_in_dit,321 overflow_module_config=dict(322 offload_dtype=dtype,323 offload_device="cpu",324 onload_dtype=dtype,325 onload_device="cpu",326 computation_dtype=self.param_dtype,327 computation_device=self.device,328 ),329 )330 self.enable_cpu_offload()331 332 def enable_cpu_offload(self):333 self.cpu_offload = True334 335 def load_models_to_device(self, loadmodel_names=[]):336 # only load models to device if cpu_offload is enabled337 if not self.cpu_offload:338 return339 # offload the unneeded models to cpu340 for model_name in self.model_names:341 if model_name not in loadmodel_names:342 model = getattr(self, model_name)343 344 if not isinstance(model, nn.Module):345 model = model.model346 347 if model is not None:348 if (349 hasattr(model, "vram_management_enabled")350 and model.vram_management_enabled351 ):352 for module in model.modules():353 if hasattr(module, "offload"):354 module.offload()355 else:356 model.cpu()357 # load the needed models to device358 for model_name in loadmodel_names:359 model = getattr(self, model_name)360 if not isinstance(model, nn.Module):361 model = model.model362 if model is not None:363 if (364 hasattr(model, "vram_management_enabled")365 and model.vram_management_enabled366 ):367 for module in model.modules():368 if hasattr(module, "onload"):369 module.onload()370 else:371 model.to(self.device)372 # fresh the cuda cache373 torch.cuda.empty_cache()374 375 376 def generate_infinitetalk(self,377 input_data,378 size_buckget='infinitetalk-480',379 motion_frame=25,380 frame_num=81,381 shift=5.0,382 sampling_steps=40,383 text_guide_scale=5.0,384 audio_guide_scale=4.0,385 n_prompt="",386 seed=-1,387 offload_model=True,388 max_frames_num=1000,389 face_scale=0.05,390 progress=True,391 color_correction_strength=0.0,392 extra_args=None):393 r"""394 Generates video frames from input image and text prompt using diffusion process.395 396 Args:397 frame_num (`int`, *optional*, defaults to 81):398 How many frames to sample from a video. The number should be 4n+1399 shift (`float`, *optional*, defaults to 5.0):400 Noise schedule shift parameter. Affects temporal dynamics401 [NOTE]: If you want to generate a 480p video, it is recommended to set the shift value to 3.0.402 sampling_steps (`int`, *optional*, defaults to 40):403 Number of diffusion sampling steps. Higher values improve quality but slow generation404 n_prompt (`str`, *optional*, defaults to ""):405 Negative prompt for content exclusion. If not given, use `config.sample_neg_prompt`406 seed (`int`, *optional*, defaults to -1):407 Random seed for noise generation. If -1, use random seed408 offload_model (`bool`, *optional*, defaults to True):409 If True, offloads models to CPU during generation to save VRAM410 """411 412 # init teacache413 if extra_args.use_teacache:414 self.model.teacache_init(415 sample_steps=sampling_steps,416 teacache_thresh=extra_args.teacache_thresh,417 model_scale=extra_args.size,418 )419 else:420 self.model.disable_teacache()421 422 input_prompt = input_data['prompt']423 cond_file_path = input_data['cond_video']424 codec = get_video_codec(cond_file_path)425 if codec == 'av1':426 output_video_path = 'tmp/' + '_input_h264.mp4'427 print(f"Converting {cond_file_path} from AV1 to H.264...")428 convert_video_to_h264(cond_file_path, output_video_path)429 print(f"Conversion complete! Saved as {output_video_path}")430 cond_file_path = output_video_path431 else:432 print("No conversion needed.")433 cond_image = extract_specific_frames(cond_file_path, 0)434 # cond_image = Image.fromarray(cond_image)435 436 437 # decide a proper size438 bucket_config_module = importlib.import_module("wan.utils.multitalk_utils")439 if size_buckget == 'infinitetalk-480':440 bucket_config = getattr(bucket_config_module, 'ASPECT_RATIO_627')441 elif size_buckget == 'infinitetalk-720':442 bucket_config = getattr(bucket_config_module, 'ASPECT_RATIO_960')443 444 src_h, src_w = cond_image.height, cond_image.width445 ratio = src_h / src_w446 closest_bucket = sorted(list(bucket_config.keys()), key=lambda x: abs(float(x)-ratio))[0]447 target_h, target_w = bucket_config[closest_bucket][0]448 cond_image = resize_and_centercrop(cond_image, (target_h, target_w))449 cond_image = cond_image / 255450 cond_image = (cond_image - 0.5) * 2 # normalization451 cond_image = cond_image.to(self.device) # 1 C 1 H W452 453 # Store the original image for color reference if strength > 0454 original_color_reference = None455 if color_correction_strength > 0.0:456 original_color_reference = cond_image.clone()457 458 459 # read audio embeddings460 audio_embedding_path_1 = input_data['cond_audio']['person1']461 if len(input_data['cond_audio']) == 1:462 HUMAN_NUMBER = 1463 audio_embedding_path_2 = None464 else:465 HUMAN_NUMBER = 2466 audio_embedding_path_2 = input_data['cond_audio']['person2']467 468 469 full_audio_embs = [] 470 audio_embedding_paths = [audio_embedding_path_1, audio_embedding_path_2]471 for human_idx in range(HUMAN_NUMBER): 472 audio_embedding_path = audio_embedding_paths[human_idx]473 if not os.path.exists(audio_embedding_path):474 continue475 full_audio_emb = torch.load(audio_embedding_path)476 if torch.isnan(full_audio_emb).any():477 continue478 if full_audio_emb.shape[0] <= frame_num:479 continue480 full_audio_embs.append(full_audio_emb) 481 482 assert len(full_audio_embs) == HUMAN_NUMBER, f"Aduio file not exists or length not satisfies frame nums."483 484 # preprocess text embedding485 if n_prompt == "":486 n_prompt = self.sample_neg_prompt487 if not self.t5_cpu:488 self.text_encoder.model.to(self.device)489 context, context_null = self.text_encoder([input_prompt, n_prompt], self.device)490 if offload_model:491 self.text_encoder.model.cpu()492 else:493 context = self.text_encoder([input_prompt], torch.device('cpu'))494 context_null = self.text_encoder([n_prompt], torch.device('cpu'))495 context = [t.to(self.device) for t in context]496 context_null = [t.to(self.device) for t in context_null]497 498 torch_gc()499 # prepare params for video generation500 indices = (torch.arange(2 * 2 + 1) - 2) * 1 501 clip_length = frame_num502 is_first_clip = True503 arrive_last_frame = False504 cur_motion_frames_num = 1505 audio_start_idx = 0506 audio_end_idx = audio_start_idx + clip_length507 gen_video_list = []508 torch_gc()509 510 # set random seed and init noise511 seed = seed if seed >= 0 else random.randint(0, 99999999)512 torch.manual_seed(seed)513 torch.cuda.manual_seed_all(seed)514 np.random.seed(seed)515 random.seed(seed)516 torch.backends.cudnn.deterministic = True517 518 # start video generation iteratively519 while True:520 audio_embs = []521 # split audio with window size522 for human_idx in range(HUMAN_NUMBER): 523 center_indices = torch.arange(524 audio_start_idx,525 audio_end_idx,526 1,527 ).unsqueeze(528 1529 ) + indices.unsqueeze(0)530 center_indices = torch.clamp(center_indices, min=0, max=full_audio_embs[human_idx].shape[0]-1)531 audio_emb = full_audio_embs[human_idx][center_indices][None,...].to(self.device)532 audio_embs.append(audio_emb)533 audio_embs = torch.concat(audio_embs, dim=0).to(self.param_dtype)534 torch_gc()535 536 h, w = cond_image.shape[-2], cond_image.shape[-1]537 lat_h, lat_w = h // self.vae_stride[1], w // self.vae_stride[2]538 max_seq_len = ((frame_num - 1) // self.vae_stride[0] + 1) * lat_h * lat_w // (539 self.patch_size[1] * self.patch_size[2])540 max_seq_len = int(math.ceil(max_seq_len / self.sp_size)) * self.sp_size541 542 543 544 noise = torch.randn(545 16, (frame_num - 1) // 4 + 1,546 lat_h,547 lat_w,548 dtype=torch.float32,549 device=self.device) 550 551 # get mask552 msk = torch.ones(1, frame_num, lat_h, lat_w, device=self.device)553 msk[:, 1:] = 0554 msk = torch.concat([555 torch.repeat_interleave(msk[:, 0:1], repeats=4, dim=1), msk[:, 1:]556 ],557 dim=1)558 msk = msk.view(1, msk.shape[1] // 4, 4, lat_h, lat_w)559 msk = msk.transpose(1, 2).to(self.param_dtype) # B 4 T H W560 561 with torch.no_grad():562 # get clip embedding563 self.clip.model.to(self.device)564 clip_context = self.clip.visual(cond_image[:, :, -1:, :, :]).to(self.param_dtype) 565 if offload_model:566 self.clip.model.cpu()567 torch_gc()568 569 # zero padding and vae encode570 video_frames = torch.zeros(1, cond_image.shape[1], frame_num-cond_image.shape[2], target_h, target_w).to(self.device)571 padding_frames_pixels_values = torch.concat([cond_image, video_frames], dim=2)572 y = self.vae.encode(padding_frames_pixels_values) 573 y = torch.stack(y).to(self.param_dtype) # B C T H W574 cur_motion_frames_latent_num = int(1 + (cur_motion_frames_num-1) // 4)575 576 if is_first_clip:577 latent_motion_frames = self.vae.encode(cond_image)[0]578 else:579 latent_motion_frames = self.vae.encode(cond_frame)[0]580 581 y = torch.concat([msk, y], dim=1) # B 4+C T H W582 torch_gc()583 584 585 # construct human mask586 human_masks = []587 if HUMAN_NUMBER==1:588 background_mask = torch.ones([src_h, src_w])589 human_mask1 = torch.ones([src_h, src_w])590 human_mask2 = torch.ones([src_h, src_w])591 human_masks = [human_mask1, human_mask2, background_mask]592 elif HUMAN_NUMBER==2:593 if 'bbox' in input_data:594 assert len(input_data['bbox']) == len(input_data['cond_audio']), f"The number of target bbox should be the same with cond_audio"595 background_mask = torch.zeros([src_h, src_w])596 for _, person_bbox in input_data['bbox'].items():597 x_min, y_min, x_max, y_max = person_bbox598 human_mask = torch.zeros([src_h, src_w])599 human_mask[int(x_min):int(x_max), int(y_min):int(y_max)] = 1600 background_mask += human_mask601 human_masks.append(human_mask)602 else:603 x_min, x_max = int(src_h * face_scale), int(src_h * (1 - face_scale))604 background_mask = torch.zeros([src_h, src_w])605 background_mask = torch.zeros([src_h, src_w])606 human_mask1 = torch.zeros([src_h, src_w])607 human_mask2 = torch.zeros([src_h, src_w])608 lefty_min, lefty_max = int((src_w//2) * face_scale), int((src_w//2) * (1 - face_scale))609 righty_min, righty_max = int((src_w//2) * face_scale + (src_w//2)), int((src_w//2) * (1 - face_scale) + (src_w//2))610 human_mask1[x_min:x_max, lefty_min:lefty_max] = 1611 human_mask2[x_min:x_max, righty_min:righty_max] = 1612 background_mask += human_mask1613 background_mask += human_mask2614 human_masks = [human_mask1, human_mask2]615 background_mask = torch.where(background_mask > 0, torch.tensor(0), torch.tensor(1))616 human_masks.append(background_mask)617 618 ref_target_masks = torch.stack(human_masks, dim=0).to(self.device)619 # resize and centercrop for ref_target_masks 620 ref_target_masks = resize_and_centercrop(ref_target_masks, (target_h, target_w))621 622 _, _, _,lat_h, lat_w = y.shape623 ref_target_masks = F.interpolate(ref_target_masks.unsqueeze(0), size=(lat_h, lat_w), mode='nearest').squeeze() 624 ref_target_masks = (ref_target_masks > 0) 625 ref_target_masks = ref_target_masks.float().to(self.device)626 627 torch_gc()628 629 @contextmanager630 def noop_no_sync():631 yield632 633 no_sync = getattr(self.model, 'no_sync', noop_no_sync)634 635 # evaluation mode636 with torch.no_grad(), no_sync():637 638 # prepare timesteps639 timesteps = list(np.linspace(self.num_timesteps, 1, sampling_steps, dtype=np.float32))640 timesteps.append(0.)641 timesteps = [torch.tensor([t], device=self.device) for t in timesteps]642 if self.use_timestep_transform:643 timesteps = [timestep_transform(t, shift=shift, num_timesteps=self.num_timesteps) for t in timesteps]644 645 # sample videos646 latent = noise647 648 # prepare condition and uncondition configs649 arg_c = {650 'context': [context],651 'clip_fea': clip_context,652 'seq_len': max_seq_len,653 'y': y,654 'audio': audio_embs,655 'ref_target_masks': ref_target_masks656 }657 658 659 arg_null_text = {660 'context': [context_null],661 'clip_fea': clip_context,662 'seq_len': max_seq_len,663 'y': y,664 'audio': audio_embs,665 'ref_target_masks': ref_target_masks666 }667 668 arg_null_audio = {669 'context': [context],670 'clip_fea': clip_context,671 'seq_len': max_seq_len,672 'y': y,673 'audio': torch.zeros_like(audio_embs)[-1:],674 'ref_target_masks': ref_target_masks675 }676 677 678 arg_null = {679 'context': [context_null],680 'clip_fea': clip_context,681 'seq_len': max_seq_len,682 'y': y,683 'audio': torch.zeros_like(audio_embs)[-1:],684 'ref_target_masks': ref_target_masks685 }686 687 torch_gc()688 if not self.vram_management:689 self.model.to(self.device)690 else:691 self.load_models_to_device(["model"])692 693 # injecting motion frames694 if not is_first_clip:695 latent_motion_frames = latent_motion_frames.to(latent.dtype).to(self.device)696 motion_add_noise = torch.randn_like(latent_motion_frames).contiguous()697 add_latent = self.add_noise(latent_motion_frames, motion_add_noise, timesteps[0])698 _, T_m, _, _ = add_latent.shape699 latent[:, :T_m] = add_latent700 701 # infer with APG702 # refer https://arxiv.org/abs/2410.02416 703 if extra_args.use_apg: 704 text_momentumbuffer = MomentumBuffer(extra_args.apg_momentum) 705 audio_momentumbuffer = MomentumBuffer(extra_args.apg_momentum) 706 707 708 progress_wrap = partial(tqdm, total=len(timesteps)-1) if progress else (lambda x: x)709 for i in progress_wrap(range(len(timesteps)-1)):710 timestep = timesteps[i]711 latent[:, :cur_motion_frames_latent_num] = latent_motion_frames712 latent_model_input = [latent.to(self.device)]713 714 # inference with CFG strategy715 noise_pred_cond = self.model(716 latent_model_input, t=timestep, **arg_c)[0] 717 torch_gc()718 719 if math.isclose(text_guide_scale, 1.0):720 noise_pred_drop_audio = self.model(721 latent_model_input, t=timestep, **arg_null_audio)[0] 722 torch_gc()723 else:724 noise_pred_drop_text = self.model(725 latent_model_input, t=timestep, **arg_null_text)[0] 726 torch_gc()727 noise_pred_uncond = self.model(728 latent_model_input, t=timestep, **arg_null)[0] 729 torch_gc()730 731 if extra_args.use_apg:732 # correct update direction733 if math.isclose(text_guide_scale, 1.0):734 diff_uncond_audio = noise_pred_cond - noise_pred_drop_audio735 noise_pred = noise_pred_cond + (audio_guide_scale - 1)* adaptive_projected_guidance(diff_uncond_audio, 736 noise_pred_cond, 737 momentum_buffer=audio_momentumbuffer, 738 norm_threshold=extra_args.apg_norm_threshold)739 else:740 diff_uncond_text = noise_pred_cond - noise_pred_drop_text741 diff_uncond_audio = noise_pred_drop_text - noise_pred_uncond742 noise_pred = noise_pred_cond + (text_guide_scale - 1) * adaptive_projected_guidance(diff_uncond_text, 743 noise_pred_cond, 744 momentum_buffer=text_momentumbuffer, 745 norm_threshold=extra_args.apg_norm_threshold) \746 + (audio_guide_scale - 1) * adaptive_projected_guidance(diff_uncond_audio, 747 noise_pred_cond, 748 momentum_buffer=audio_momentumbuffer, 749 norm_threshold=extra_args.apg_norm_threshold)750 else:751 # vanilla CFG strategy752 if math.isclose(text_guide_scale, 1.0):753 noise_pred = noise_pred_drop_audio + audio_guide_scale* (noise_pred_cond - noise_pred_drop_audio) 754 else:755 noise_pred = noise_pred_uncond + text_guide_scale * (756 noise_pred_cond - noise_pred_drop_text) + \757 audio_guide_scale * (noise_pred_drop_text - noise_pred_uncond) 758 noise_pred = -noise_pred 759 760 # update latent761 dt = timesteps[i] - timesteps[i + 1]762 dt = dt / self.num_timesteps763 latent = latent + noise_pred * dt[:, None, None, None]764 765 # injecting motion frames766 if not is_first_clip:767 latent_motion_frames = latent_motion_frames.to(latent.dtype).to(self.device)768 motion_add_noise = torch.randn_like(latent_motion_frames).contiguous()769 add_latent = self.add_noise(latent_motion_frames, motion_add_noise, timesteps[i+1])770 _, T_m, _, _ = add_latent.shape771 latent[:, :T_m] = add_latent772 773 latent[:, :cur_motion_frames_latent_num] = latent_motion_frames774 x0 = [latent.to(self.device)] 775 del latent_model_input, timestep776 777 if offload_model: 778 if not self.vram_management:779 self.model.cpu()780 torch_gc()781 782 videos = self.vae.decode(x0)783 784 # cache generated samples785 videos = torch.stack(videos).cpu() # B C T H W786 # >>> START OF COLOR CORRECTION STEP <<<787 if color_correction_strength > 0.0 and original_color_reference is not None:788 videos = match_and_blend_colors(videos, original_color_reference, color_correction_strength)789 # >>> END OF COLOR CORRECTION STEP <<<790 791 if is_first_clip:792 gen_video_list.append(videos)793 else:794 gen_video_list.append(videos[:, :, cur_motion_frames_num:])795 796 # decide whether is done797 if arrive_last_frame: break798 799 # update next condition frames800 is_first_clip = False801 cur_motion_frames_num = motion_frame802 803 cond_frame = videos[:, :, -cur_motion_frames_num:].to(torch.float32).to(self.device)804 audio_start_idx += (frame_num - cur_motion_frames_num)805 audio_end_idx = audio_start_idx + clip_length806 807 cond_image = extract_specific_frames(cond_file_path, audio_start_idx)808 # cond_image = Image.fromarray(cond_image)809 cond_image = resize_and_centercrop(cond_image, (target_h, target_w))810 cond_image = cond_image / 255811 cond_image = (cond_image - 0.5) * 2 # normalization812 cond_image = cond_image.to(self.device) # 1 C 1 H W813 814 # Repeat audio emb815 if audio_end_idx >= min(max_frames_num, len(full_audio_embs[0])):816 arrive_last_frame = True817 miss_lengths = []818 source_frames = []819 for human_inx in range(HUMAN_NUMBER):820 source_frame = len(full_audio_embs[human_inx])821 source_frames.append(source_frame)822 if audio_end_idx >= len(full_audio_embs[human_inx]):823 miss_length = audio_end_idx - len(full_audio_embs[human_inx]) + 3 824 add_audio_emb = torch.flip(full_audio_embs[human_inx][-1*miss_length:], dims=[0])825 full_audio_embs[human_inx] = torch.cat([full_audio_embs[human_inx], add_audio_emb], dim=0)826 miss_lengths.append(miss_length)827 else:828 miss_lengths.append(0)829 830 831 if max_frames_num <= frame_num: break832 833 torch_gc()834 if offload_model: 835 torch.cuda.synchronize()836 if dist.is_initialized():837 dist.barrier()838 839 gen_video_samples = torch.cat(gen_video_list, dim=2)[:, :, :int(max_frames_num)] 840 gen_video_samples = gen_video_samples.to(torch.float32)841 if max_frames_num > frame_num and sum(miss_lengths) > 0:842 # split video frames843 # gen_video_samples = gen_video_samples[:, :, :-1*miss_lengths[0]]844 gen_video_samples = gen_video_samples[:, :, :full_audio_emb.shape[0]]845 846 if dist.is_initialized():847 dist.barrier()848 849 del noise, latent850 torch_gc()851 852 return gen_video_samples[0] if self.rank == 0 else None853 854 855 856