durgappc/infinitetalk2

1# Copyright 2024-2025 The Alibaba Wan Team Authors. All rights reserved.2import argparse3import binascii4import os5import os.path as osp6import cv27 8import imageio9import torch10import torchvision11from PIL import Image12import librosa13import soundfile as sf14import subprocess15from decord import VideoReader, cpu16import gc17 18__all__ = ['cache_video', 'cache_image', 'str2bool']19 20 21def rand_name(length=8, suffix=''):22    name = binascii.b2a_hex(os.urandom(length)).decode('utf-8')23    if suffix:24        if not suffix.startswith('.'):25            suffix = '.' + suffix26        name += suffix27    return name28 29 30 31def str2bool(v):32    """33    Convert a string to a boolean.34 35    Supported true values: 'yes', 'true', 't', 'y', '1'36    Supported false values: 'no', 'false', 'f', 'n', '0'37 38    Args:39        v (str): String to convert.40 41    Returns:42        bool: Converted boolean value.43 44    Raises:45        argparse.ArgumentTypeError: If the value cannot be converted to boolean.46    """47    if isinstance(v, bool):48        return v49    v_lower = v.lower()50    if v_lower in ('yes', 'true', 't', 'y', '1'):51        return True52    elif v_lower in ('no', 'false', 'f', 'n', '0'):53        return False54    else:55        raise argparse.ArgumentTypeError('Boolean value expected (True/False)')56    57def cache_video(tensor,58                save_file=None,59                fps=30,60                suffix='.mp4',61                nrow=8,62                normalize=True,63                value_range=(-1, 1),64                retry=5):65    # cache file66    cache_file = osp.join('/tmp', rand_name(67        suffix=suffix)) if save_file is None else save_file68 69    # save to cache70    error = None71    for _ in range(retry):72        try:73            # preprocess74            tensor = tensor.clamp(min(value_range), max(value_range))75            tensor = torch.stack([76                torchvision.utils.make_grid(77                    u, nrow=nrow, normalize=normalize, value_range=value_range)78                for u in tensor.unbind(2)79            ],80                                 dim=1).permute(1, 2, 3, 0)81            tensor = (tensor * 255).type(torch.uint8).cpu()82 83            # write video84            writer = imageio.get_writer(85                cache_file, fps=fps, codec='libx264', quality=8)86            for frame in tensor.numpy():87                writer.append_data(frame)88            writer.close()89            return cache_file90        except Exception as e:91            error = e92            continue93    else:94        print(f'cache_video failed, error: {error}', flush=True)95        return None96 97 98def cache_image(tensor,99                save_file,100                nrow=8,101                normalize=True,102                value_range=(-1, 1),103                retry=5):104    # cache file105    suffix = osp.splitext(save_file)[1]106    if suffix.lower() not in [107            '.jpg', '.jpeg', '.png', '.tiff', '.gif', '.webp'108    ]:109        suffix = '.png'110 111    # save to cache112    error = None113    for _ in range(retry):114        try:115            tensor = tensor.clamp(min(value_range), max(value_range))116            torchvision.utils.save_image(117                tensor,118                save_file,119                nrow=nrow,120                normalize=normalize,121                value_range=value_range)122            return save_file123        except Exception as e:124            error = e125            continue126 127def convert_video_to_h264(input_video_path, output_video_path):128    subprocess.run(129        ['ffmpeg', '-i', input_video_path, '-c:v', 'libx264', '-c:a', 'copy', output_video_path],130        stdout=subprocess.PIPE,131        stderr=subprocess.PIPE132    )133 134 135def is_video(path):136    video_exts = ['.mp4', '.avi', '.mov', '.mkv', '.flv', '.wmv', '.webm', '.mpeg', '.mpg']137    return os.path.splitext(path)[1].lower() in video_exts138 139 140def extract_specific_frames(video_path, frame_id):141    if is_video(video_path):142        vr = VideoReader(video_path, ctx=cpu(0))143        if frame_id < vr._num_frame:144            frame = vr[frame_id].asnumpy()  # RGB145        else:146            frame = vr[-1].asnumpy()147        del vr148        gc.collect()149        frame = Image.fromarray(frame)150    else:151        frame = Image.open(video_path).convert("RGB")152    return frame153 154def get_video_codec(video_path):155    result = subprocess.run(156        ['ffprobe', '-v', 'error', '-select_streams', 'v:0',157         '-show_entries', 'stream=codec_name', '-of', 'default=nw=1:nk=1', video_path],158        stdout=subprocess.PIPE,159        stderr=subprocess.PIPE160    )161    codec = result.stdout.decode().strip()162    return codec163 164 165 166def split_wav_librosa(wav_path, segments, save_dir):167    y, sr = librosa.load(wav_path, sr=None)168    filename = wav_path.split('/')[-1].split('.')[0]169    save_list = []170    for idx, (start, end) in enumerate(segments):171        start_sample = int(start * sr)172        end_sample = int(end * sr)173        segment = y[start_sample:end_sample]174        out_path = os.path.join(save_dir, filename + str(start) + '_' + str(end) + '.wav')175        sf.write(out_path, segment, sr)176        print(f"Saved {out_path}: {start}s to {end}s")177        save_list.append(out_path)178    return save_list179 180