DFAGWE/infinitetalk2

1import os2from einops import rearrange3 4import torch5import torch.nn as nn6 7from xfuser.core.distributed import (8    get_sequence_parallel_rank,9    get_sequence_parallel_world_size,10    get_sp_group,11)12from einops import rearrange, repeat13from functools import lru_cache14import imageio15import uuid16from tqdm import tqdm17import numpy as np18import subprocess19import soundfile as sf20import torchvision21import binascii22import os.path as osp23from skimage import color24 25VID_EXTENSIONS = (".mp4", ".avi", ".mov", ".mkv")26ASPECT_RATIO_627 = {27     '0.26': ([320, 1216], 1), '0.38': ([384, 1024], 1), '0.50': ([448, 896], 1), '0.67': ([512, 768], 1), 28     '0.82': ([576, 704], 1),  '1.00': ([640, 640], 1),  '1.22': ([704, 576], 1), '1.50': ([768, 512], 1), 29     '1.86': ([832, 448], 1),  '2.00': ([896, 448], 1),  '2.50': ([960, 384], 1), '2.83': ([1088, 384], 1), 30     '3.60': ([1152, 320], 1), '3.80': ([1216, 320], 1), '4.00': ([1280, 320], 1)}31 32 33ASPECT_RATIO_960 = {34     '0.22': ([448, 2048], 1), '0.29': ([512, 1792], 1), '0.36': ([576, 1600], 1), '0.45': ([640, 1408], 1), 35     '0.55': ([704, 1280], 1), '0.63': ([768, 1216], 1), '0.76': ([832, 1088], 1), '0.88': ([896, 1024], 1), 36     '1.00': ([960, 960], 1), '1.14': ([1024, 896], 1), '1.31': ([1088, 832], 1), '1.50': ([1152, 768], 1), 37     '1.58': ([1216, 768], 1), '1.82': ([1280, 704], 1), '1.91': ([1344, 704], 1), '2.20': ([1408, 640], 1), 38     '2.30': ([1472, 640], 1), '2.67': ([1536, 576], 1), '2.89': ([1664, 576], 1), '3.62': ([1856, 512], 1), 39     '3.75': ([1920, 512], 1)}40 41 42 43def torch_gc():44    torch.cuda.empty_cache()45    torch.cuda.ipc_collect()46 47 48 49def split_token_counts_and_frame_ids(T, token_frame, world_size, rank):50 51    S = T * token_frame52    split_sizes = [S // world_size + (1 if i < S % world_size else 0) for i in range(world_size)]53    start = sum(split_sizes[:rank])54    end = start + split_sizes[rank]55    counts = [0] * T56    for idx in range(start, end):57        t = idx // token_frame58        counts[t] += 159 60    counts_filtered = []61    frame_ids = []62    for t, c in enumerate(counts):63        if c > 0:64            counts_filtered.append(c)65            frame_ids.append(t)66    return counts_filtered, frame_ids67 68 69def normalize_and_scale(column, source_range, target_range, epsilon=1e-8):70 71    source_min, source_max = source_range72    new_min, new_max = target_range73 74    normalized = (column - source_min) / (source_max - source_min + epsilon)75    scaled = normalized * (new_max - new_min) + new_min76    return scaled77 78 79@torch.compile80def calculate_x_ref_attn_map(visual_q, ref_k, ref_target_masks, mode='mean', attn_bias=None):81    82    ref_k = ref_k.to(visual_q.dtype).to(visual_q.device)83    scale = 1.0 / visual_q.shape[-1] ** 0.584    visual_q = visual_q * scale85    visual_q = visual_q.transpose(1, 2)86    ref_k = ref_k.transpose(1, 2)87    attn = visual_q @ ref_k.transpose(-2, -1)88 89    if attn_bias is not None:90        attn = attn + attn_bias91 92    x_ref_attn_map_source = attn.softmax(-1) # B, H, x_seqlens, ref_seqlens93 94 95    x_ref_attn_maps = []96    ref_target_masks = ref_target_masks.to(visual_q.dtype)97    x_ref_attn_map_source = x_ref_attn_map_source.to(visual_q.dtype)98 99    for class_idx, ref_target_mask in enumerate(ref_target_masks):100        torch_gc()101        ref_target_mask = ref_target_mask[None, None, None, ...]102        x_ref_attnmap = x_ref_attn_map_source * ref_target_mask103        x_ref_attnmap = x_ref_attnmap.sum(-1) / ref_target_mask.sum() # B, H, x_seqlens, ref_seqlens --> B, H, x_seqlens104        x_ref_attnmap = x_ref_attnmap.permute(0, 2, 1) # B, x_seqlens, H105       106        if mode == 'mean':107            x_ref_attnmap = x_ref_attnmap.mean(-1) # B, x_seqlens108        elif mode == 'max':109            x_ref_attnmap = x_ref_attnmap.max(-1) # B, x_seqlens110        111        x_ref_attn_maps.append(x_ref_attnmap)112    113    del attn114    del x_ref_attn_map_source115    torch_gc()116 117    return torch.concat(x_ref_attn_maps, dim=0)118 119 120def get_attn_map_with_target(visual_q, ref_k, shape, ref_target_masks=None, split_num=2, enable_sp=False):121    """Args:122        query (torch.tensor): B M H K123        key (torch.tensor): B M H K124        shape (tuple): (N_t, N_h, N_w)125        ref_target_masks: [B, N_h * N_w]126    """127 128    N_t, N_h, N_w = shape129    if enable_sp:130        ref_k = get_sp_group().all_gather(ref_k, dim=1)131    132    x_seqlens = N_h * N_w133    ref_k     = ref_k[:, :x_seqlens]134    _, seq_lens, heads, _ = visual_q.shape135    class_num, _ = ref_target_masks.shape136    x_ref_attn_maps = torch.zeros(class_num, seq_lens).to(visual_q.device).to(visual_q.dtype)137 138    split_chunk = heads // split_num139    140    for i in range(split_num):141        x_ref_attn_maps_perhead = calculate_x_ref_attn_map(visual_q[:, :, i*split_chunk:(i+1)*split_chunk, :], ref_k[:, :, i*split_chunk:(i+1)*split_chunk, :], ref_target_masks)142        x_ref_attn_maps += x_ref_attn_maps_perhead143    144    return x_ref_attn_maps / split_num145 146 147def rotate_half(x):148    x = rearrange(x, "... (d r) -> ... d r", r=2)149    x1, x2 = x.unbind(dim=-1)150    x = torch.stack((-x2, x1), dim=-1)151    return rearrange(x, "... d r -> ... (d r)")152 153 154class RotaryPositionalEmbedding1D(nn.Module):155 156    def __init__(self,157                 head_dim,158                 ):159        super().__init__()160        self.head_dim = head_dim161        self.base = 10000162 163 164    @lru_cache(maxsize=32)165    def precompute_freqs_cis_1d(self, pos_indices):166 167        freqs = 1.0 / (self.base ** (torch.arange(0, self.head_dim, 2)[: (self.head_dim // 2)].float() / self.head_dim))168        freqs = freqs.to(pos_indices.device)169        freqs = torch.einsum("..., f -> ... f", pos_indices.float(), freqs)170        freqs = repeat(freqs, "... n -> ... (n r)", r=2)171        return freqs172 173    def forward(self, x, pos_indices):174        """1D RoPE.175 176        Args:177            query (torch.tensor): [B, head, seq, head_dim]178            pos_indices (torch.tensor): [seq,]179        Returns:180            query with the same shape as input.181        """182        freqs_cis = self.precompute_freqs_cis_1d(pos_indices)183 184        x_ = x.float()185 186        freqs_cis = freqs_cis.float().to(x.device)187        cos, sin = freqs_cis.cos(), freqs_cis.sin()188        cos, sin = rearrange(cos, 'n d -> 1 1 n d'), rearrange(sin, 'n d -> 1 1 n d')189        x_ = (x_ * cos) + (rotate_half(x_) * sin)190 191        return x_.type_as(x)192    193 194 195def rand_name(length=8, suffix=''):196    name = binascii.b2a_hex(os.urandom(length)).decode('utf-8')197    if suffix:198        if not suffix.startswith('.'):199            suffix = '.' + suffix200        name += suffix201    return name202 203def cache_video(tensor,204                save_file=None,205                fps=30,206                suffix='.mp4',207                nrow=8,208                normalize=True,209                value_range=(-1, 1),210                retry=5):211    212    # cache file213    cache_file = osp.join('/tmp', rand_name(214        suffix=suffix)) if save_file is None else save_file215 216    # save to cache217    error = None218    for _ in range(retry):219       220        # preprocess221        tensor = tensor.clamp(min(value_range), max(value_range))222        tensor = torch.stack([223                torchvision.utils.make_grid(224                    u, nrow=nrow, normalize=normalize, value_range=value_range)225                for u in tensor.unbind(2)226            ],227                                 dim=1).permute(1, 2, 3, 0)228        tensor = (tensor * 255).type(torch.uint8).cpu()229 230        # write video231        writer = imageio.get_writer(cache_file, fps=fps, codec='libx264', quality=10, ffmpeg_params=["-crf", "10"])232        for frame in tensor.numpy():233            writer.append_data(frame)234        writer.close()235        return cache_file236 237def save_video_ffmpeg(gen_video_samples, save_path, vocal_audio_list, fps=25, quality=5, high_quality_save=False):238    239    def save_video(frames, save_path, fps, quality=9, ffmpeg_params=None):240        writer = imageio.get_writer(241            save_path, fps=fps, quality=quality, ffmpeg_params=ffmpeg_params242        )243        for frame in tqdm(frames, desc="Saving video"):244            frame = np.array(frame)245            writer.append_data(frame)246        writer.close()247    save_path_tmp = save_path + "-temp.mp4"248 249    if high_quality_save:250        cache_video(251                    tensor=gen_video_samples.unsqueeze(0),252                    save_file=save_path_tmp,253                    fps=fps,254                    nrow=1,255                    normalize=True,256                    value_range=(-1, 1)257                    )258    else:259        video_audio = (gen_video_samples+1)/2 # C T H W260        video_audio = video_audio.permute(1, 2, 3, 0).cpu().numpy()261        video_audio = np.clip(video_audio * 255, 0, 255).astype(np.uint8)  # to [0, 255]262        save_video(video_audio, save_path_tmp, fps=fps, quality=quality)263 264 265    # crop audio according to video length266    _, T, _, _ = gen_video_samples.shape267    duration = T / fps268    save_path_crop_audio = save_path + "-cropaudio.wav"269    final_command = [270        "ffmpeg",271        "-i",272        vocal_audio_list[0],273        "-t",274        f'{duration}',275        save_path_crop_audio,276    ]277    subprocess.run(final_command, check=True)278 279    save_path = save_path + ".mp4"280    if high_quality_save:281        final_command = [282            "ffmpeg",283            "-y",284            "-i", save_path_tmp,285            "-i", save_path_crop_audio,286            "-c:v", "libx264",287            "-crf", "0",288            "-preset", "veryslow",289            "-c:a", "aac", 290            "-shortest",291            save_path,292        ]293        subprocess.run(final_command, check=True)294        os.remove(save_path_tmp)295        os.remove(save_path_crop_audio)296    else:297        final_command = [298            "ffmpeg",299            "-y",300            "-i",301            save_path_tmp,302            "-i",303            save_path_crop_audio,304            "-c:v",305            "libx264",306            "-c:a",307            "aac",308            "-shortest",309            save_path,310        ]311        subprocess.run(final_command, check=True)312        os.remove(save_path_tmp)313        os.remove(save_path_crop_audio)314 315 316class MomentumBuffer:317    def __init__(self, momentum: float): 318        self.momentum = momentum 319        self.running_average = 0 320    321    def update(self, update_value: torch.Tensor): 322        new_average = self.momentum * self.running_average 323        self.running_average = update_value + new_average324    325 326 327def project( 328        v0: torch.Tensor, # [B, C, T, H, W] 329        v1: torch.Tensor, # [B, C, T, H, W] 330        ): 331    dtype = v0.dtype 332    v0, v1 = v0.double(), v1.double() 333    v1 = torch.nn.functional.normalize(v1, dim=[-1, -2, -3, -4]) 334    v0_parallel = (v0 * v1).sum(dim=[-1, -2, -3, -4], keepdim=True) * v1 335    v0_orthogonal = v0 - v0_parallel336    return v0_parallel.to(dtype), v0_orthogonal.to(dtype)337 338 339def adaptive_projected_guidance( 340          diff: torch.Tensor, # [B, C, T, H, W] 341          pred_cond: torch.Tensor, # [B, C, T, H, W] 342          momentum_buffer: MomentumBuffer = None, 343          eta: float = 0.0,344          norm_threshold: float = 55,345          ): 346    if momentum_buffer is not None: 347        momentum_buffer.update(diff) 348        diff = momentum_buffer.running_average349    if norm_threshold > 0: 350        ones = torch.ones_like(diff) 351        diff_norm = diff.norm(p=2, dim=[-1, -2, -3, -4], keepdim=True) 352        print(f"diff_norm: {diff_norm}")353        scale_factor = torch.minimum(ones, norm_threshold / diff_norm) 354        diff = diff * scale_factor 355    diff_parallel, diff_orthogonal = project(diff, pred_cond) 356    normalized_update = diff_orthogonal + eta * diff_parallel357    return normalized_update358 359 360 361def match_and_blend_colors(source_chunk: torch.Tensor, reference_image: torch.Tensor, strength: float) -> torch.Tensor:362    """363    Matches the color of a source video chunk to a reference image and blends with the original.364 365    Args:366        source_chunk (torch.Tensor): The video chunk to be color-corrected (B, C, T, H, W) in range [-1, 1].367                                     Assumes B=1 (batch size of 1).368        reference_image (torch.Tensor): The reference image (B, C, 1, H, W) in range [-1, 1].369                                        Assumes B=1 and T=1 (single reference frame).370        strength (float): The strength of the color correction (0.0 to 1.0).371                          0.0 means no correction, 1.0 means full correction.372 373    Returns:374        torch.Tensor: The color-corrected and blended video chunk.375    """376    # print(f"[match_and_blend_colors] Input source_chunk shape: {source_chunk.shape}, reference_image shape: {reference_image.shape}, strength: {strength}")377 378    if strength == 0.0:379        # print(f"[match_and_blend_colors] Strength is 0, returning original source_chunk.")380        return source_chunk381 382    if not 0.0 <= strength <= 1.0:383        raise ValueError(f"Strength must be between 0.0 and 1.0, got {strength}")384 385    device = source_chunk.device386    dtype = source_chunk.dtype387 388    # Squeeze batch dimension, permute to T, H, W, C for skimage389    # Source: (1, C, T, H, W) -> (T, H, W, C)390    source_np = source_chunk.squeeze(0).permute(1, 2, 3, 0).cpu().numpy()391    # Reference: (1, C, 1, H, W) -> (H, W, C)392    ref_np = reference_image.squeeze(0).squeeze(1).permute(1, 2, 0).cpu().numpy() # Squeeze T dimension as well393 394    # Normalize from [-1, 1] to [0, 1] for skimage395    source_np_01 = (source_np + 1.0) / 2.0396    ref_np_01 = (ref_np + 1.0) / 2.0397 398    # Clip to ensure values are strictly in [0, 1] after potential float precision issues399    source_np_01 = np.clip(source_np_01, 0.0, 1.0)400    ref_np_01 = np.clip(ref_np_01, 0.0, 1.0)401 402    # Convert reference to Lab403    try:404        ref_lab = color.rgb2lab(ref_np_01)405    except ValueError as e:406        # Handle potential errors if image data is not valid for conversion407        print(f"Warning: Could not convert reference image to Lab: {e}. Skipping color correction for this chunk.")408        return source_chunk409 410 411    corrected_frames_np_01 = []412    for i in range(source_np_01.shape[0]): # Iterate over time (T)413        source_frame_rgb_01 = source_np_01[i]414        415        try:416            source_lab = color.rgb2lab(source_frame_rgb_01)417        except ValueError as e:418            print(f"Warning: Could not convert source frame {i} to Lab: {e}. Using original frame.")419            corrected_frames_np_01.append(source_frame_rgb_01)420            continue421 422        corrected_lab_frame = source_lab.copy()423 424        # Perform color transfer for L, a, b channels425        for j in range(3): # L, a, b426            mean_src, std_src = source_lab[:, :, j].mean(), source_lab[:, :, j].std()427            mean_ref, std_ref = ref_lab[:, :, j].mean(), ref_lab[:, :, j].std()428 429            # Avoid division by zero if std_src is 0430            if std_src == 0:431                # If source channel has no variation, keep it as is, but shift by reference mean432                # This case is debatable, could also just copy source or target mean.433                # Shifting by target mean helps if source is flat but target isn't.434                corrected_lab_frame[:, :, j] = mean_ref 435            else:436                corrected_lab_frame[:, :, j] = (corrected_lab_frame[:, :, j] - mean_src) * (std_ref / std_src) + mean_ref437        438        try:439            fully_corrected_frame_rgb_01 = color.lab2rgb(corrected_lab_frame)440        except ValueError as e:441            print(f"Warning: Could not convert corrected frame {i} back to RGB: {e}. Using original frame.")442            corrected_frames_np_01.append(source_frame_rgb_01)443            continue444            445        # Clip again after lab2rgb as it can go slightly out of [0,1]446        fully_corrected_frame_rgb_01 = np.clip(fully_corrected_frame_rgb_01, 0.0, 1.0)447 448        # Blend with original source frame (in [0,1] RGB)449        blended_frame_rgb_01 = (1 - strength) * source_frame_rgb_01 + strength * fully_corrected_frame_rgb_01450        corrected_frames_np_01.append(blended_frame_rgb_01)451 452    corrected_chunk_np_01 = np.stack(corrected_frames_np_01, axis=0)453 454    # Convert back to [-1, 1]455    corrected_chunk_np_minus1_1 = (corrected_chunk_np_01 * 2.0) - 1.0456 457    # Permute back to (C, T, H, W), add batch dim, and convert to original torch.Tensor type and device458    # (T, H, W, C) -> (C, T, H, W)459    corrected_chunk_tensor = torch.from_numpy(corrected_chunk_np_minus1_1).permute(3, 0, 1, 2).unsqueeze(0)460    corrected_chunk_tensor = corrected_chunk_tensor.contiguous() # Ensure contiguous memory layout461    output_tensor = corrected_chunk_tensor.to(device=device, dtype=dtype)462    # print(f"[match_and_blend_colors] Output tensor shape: {output_tensor.shape}")463    return output_tensor464