utils/training_utils.py

import os
import torch
import torch.distributed as dist
from torch.utils.data import DistributedSampler
from safetensors.torch import load_file
import random
import glob
import numpy as np
from torchvision.transforms.functional import resize
from torchvision.transforms import InterpolationMode
from torchvision.utils import save_image


def init_distributed(args):
    """初始化分布式训练环境"""
    if not dist.is_initialized():
        dist.init_process_group(backend="nccl")
    global_rank = dist.get_rank()
    world_size = dist.get_world_size()
    local_rank = int(os.environ["LOCAL_RANK"])
    torch.cuda.set_device(local_rank)

    sp_size = 1
    dp_size = world_size

    if args.use_usp:
        sp_size = args.ulysses_degree * args.ring_degree
        dp_size = world_size // sp_size
        assert sp_size <= world_size, f"sequence parallel size ({sp_size}) must be less than or equal to world size ({world_size})."
        assert world_size % sp_size == 0, f"world size ({world_size}) must be divisible by sequence parallel size ({sp_size})."
        
        from xfuser.core.distributed import initialize_model_parallel, init_distributed_environment, get_data_parallel_rank
        init_distributed_environment(rank=global_rank, world_size=world_size)
        initialize_model_parallel(
            data_parallel_degree=dp_size,
            sequence_parallel_degree=sp_size,
            ring_degree=args.ring_degree,
            ulysses_degree=args.ulysses_degree,
        )

        if dist.is_initialized():
            sp_group_id = get_data_parallel_rank()
            if global_rank == 0:
                base_seed = torch.randint(0, 1000000, (1,)).item()
            else:
                base_seed = 0
            base_seed = torch.tensor([base_seed], device="cuda")
            dist.broadcast(base_seed, src=0)
            base_seed = base_seed.item()
            seed = base_seed + sp_group_id
            torch.manual_seed(seed)

    return global_rank, world_size, local_rank, sp_size, dp_size


def custom_collate_fn(batch):
    """自定义数据整理函数，用于处理不同长度的视频"""
    keys = batch[0].keys()
    data = {}
    batch_video_length = min([item["video"].shape[1] for item in batch])
    for key in keys:
        if key=='video': # 截断视频长度
            data[key] = [item[key][:,:batch_video_length] for item in batch]
        elif key == 'audio_emb':
            data[key] = [item[key][:batch_video_length] for item in batch]
        elif key == 'audio_embed_speaker1':
            data[key] = [item[key][:batch_video_length] for item in batch]
        elif key == 'audio_embed_speaker2':
            data[key] = [item[key][:batch_video_length] for item in batch]
        else:
            data[key] = [item[key] for item in batch]
    return data


def get_distributed_sampler(dataset, args, global_rank, world_size, sp_size, dp_size):
    """获取分布式采样器"""
    if dp_size == 1:
        sampler = None
    else:
        if args.use_usp:
            from xfuser.core.distributed import get_data_parallel_rank
            dp_rank = get_data_parallel_rank()
        else:
            dp_rank = global_rank  # 避免调用 _DP 相关内容

        sampler = DistributedSampler(
            dataset,
            num_replicas=dp_size,
            rank=dp_rank,
            shuffle=False
        )

        if global_rank == 0:
            print(f"Using DistributedSampler: dp_size={dp_size}, dp_rank={dp_rank}")

    return sampler


def load_predefined_prompt_embeddings(base_dir="./multi_person/prompt_emb_concat", text_len=512):
    """
    加载三种预定义的prompt embeddings
    
    Args:
        base_dir: 预定义embedding的基础目录
        text_len: 文本长度，用于padding或截断
    
    Returns:
        dict: 包含三种类型embedding的字典
        {
            'talk_prompts': list of tensors,
            'silent_prompts_left': list of tensors, 
            'silent_prompts_right': list of tensors
        }
    """
    prompt_embeddings = {
        'talk_prompts': [],
        'silent_prompts_left': [],
        'silent_prompts_right': []
    }
    
    # 定义三个子目录
    subdirs = {
        'talk_prompts': 'talk_prompts',
        'silent_prompts_left': 'silent_prompts_left', 
        'silent_prompts_right': 'silent_prompts_right'
    }
    
    for key, subdir in subdirs.items():
        dir_path = os.path.join(base_dir, subdir)
        if not os.path.exists(dir_path):
            print(f"Warning: Directory {dir_path} does not exist")
            continue
            
        # 获取所有.safetensors文件
        pattern = os.path.join(dir_path, "*.safetensors")
        files = sorted(glob.glob(pattern))
        
        for file_path in files:
            try:
                # 加载embedding
                prompt_data = load_file(file_path)
                prompt_emb = prompt_data['context']
                
                # 处理长度
                if prompt_emb.shape[0] < text_len:
                    padding = torch.zeros(text_len - prompt_emb.shape[0], prompt_emb.shape[1])
                    prompt_emb = torch.cat([prompt_emb, padding], dim=0)
                else:
                    prompt_emb = prompt_emb[:text_len]
                
                prompt_embeddings[key].append(prompt_emb)
                
            except Exception as e:
                print(f"Error loading {file_path}: {e}")
                continue
    
    print(f"Loaded {len(prompt_embeddings['talk_prompts'])} talk prompts")
    print(f"Loaded {len(prompt_embeddings['silent_prompts_left'])} silent left prompts") 
    print(f"Loaded {len(prompt_embeddings['silent_prompts_right'])} silent right prompts")
    
    return prompt_embeddings


def get_random_prompt_embedding(prompt_embeddings, prompt_type=None, device=None, dtype=None):
    """
    从预定义的prompt embeddings中随机选择一个
    
    Args:
        prompt_embeddings: 由load_predefined_prompt_embeddings返回的字典
        prompt_type: 指定类型 ('talk_prompts', 'silent_prompts_left', 'silent_prompts_right')
                    如果为None，则随机选择类型
        device: 目标设备
        dtype: 目标数据类型
    
    Returns:
        torch.Tensor: 随机选择的prompt embedding
    """
    if prompt_type is None:
        # 随机选择类型
        available_types = [k for k, v in prompt_embeddings.items() if len(v) > 0]
        if not available_types:
            raise ValueError("No prompt embeddings available")
        prompt_type = random.choice(available_types)
    
    if prompt_type not in prompt_embeddings:
        raise ValueError(f"Unknown prompt type: {prompt_type}")
    
    if len(prompt_embeddings[prompt_type]) == 0:
        raise ValueError(f"No embeddings available for type: {prompt_type}")
    
    # 随机选择一个embedding
    selected_embedding = random.choice(prompt_embeddings[prompt_type])
    
    # 移动到指定设备和类型
    if device is not None:
        selected_embedding = selected_embedding.to(device)
    if dtype is not None:
        selected_embedding = selected_embedding.to(dtype)
    
    return selected_embedding


def create_silence_video(video_tensor, cycle_frames=3):
    """
    创建沉默视频：使用前N帧循环播放，支持倒序播放模式
    
    Args:
        video_tensor: [b, 3, F, H, W] 原始视频tensor
        cycle_frames: int, 循环播放的帧数，默认为2
        
    Returns:
        silence_video: [b, 3, F, H, W] 沉默视频tensor
        
    播放模式：
    - cycle_frames=2: 12121212... (原逻辑)
    - cycle_frames=3: 123123123...
    - cycle_frames=4: 1234543212345... (倒序模式)
    - cycle_frames=5: 1234543212345...
    """
    batch_size, channels, num_frames, height, width = video_tensor.shape
    
    # 确保cycle_frames不超过视频总帧数
    cycle_frames = min(cycle_frames, num_frames)
    
    # 获取前cycle_frames帧
    cycle_video = video_tensor[:, :, :cycle_frames, :, :]  # [b, 3, cycle_frames, H, W]
    
    if cycle_frames <= 2:
        # 对于1-2帧，使用简单的交替模式
        if cycle_frames == 1:
            # 单帧重复
            silence_video = cycle_video.repeat(1, 1, num_frames, 1, 1)
        else:
            # 双帧交替：12121212...
            frame_1 = cycle_video[:, :, 0:1, :, :]  # [b, 3, 1, H, W]
            frame_2 = cycle_video[:, :, 1:2, :, :]  # [b, 3, 1, H, W]
            
            repeat_times = (num_frames + 1) // 2
            frame_1_repeated = frame_1.repeat(1, 1, repeat_times, 1, 1)
            frame_2_repeated = frame_2.repeat(1, 1, repeat_times, 1, 1)
            
            silence_frames = torch.stack([frame_1_repeated, frame_2_repeated], dim=3)
            silence_frames = silence_frames.view(batch_size, channels, repeat_times * 2, height, width)
            silence_video = silence_frames[:, :, :num_frames, :, :]
    else:
        # 对于3帧以上，使用倒序模式：1234543212345...
        # 创建一个完整的循环周期：123...cycle_frames...321
        forward_frames = cycle_video  # [b, 3, cycle_frames, H, W]
        reverse_frames = torch.flip(cycle_video[:, :, 1:-1, :, :], dims=[2])  # [b, 3, cycle_frames-2, H, W]
        
        # 拼接一个完整周期：123...cycle_frames...321
        one_cycle = torch.cat([forward_frames, reverse_frames], dim=2)  # [b, 3, 2*cycle_frames-2, H, W]
        cycle_length = 2 * cycle_frames - 2
        
        # 计算需要多少个完整周期
        num_cycles = (num_frames + cycle_length - 1) // cycle_length
        
        # 重复完整周期
        repeated_cycles = one_cycle.repeat(1, 1, num_cycles, 1, 1)  # [b, 3, num_cycles*cycle_length, H, W]
        
        # 截断到所需帧数
        silence_video = repeated_cycles[:, :, :num_frames, :, :]  # [b, 3, F, H, W]
    
    return silence_video


def extract_square_faces_from_ref_images(face_parser, ref_images, video_paths, first_frame_faces, crop_size=224, device=None, torch_dtype=None, global_rank=0, current_global_step=0):
    """
    使用人脸解析器从参考图像中提取方形人脸，如果检测不到人脸则回退到dataset提供的crop人脸
    
    Args:
        face_parser: FaceInference实例
        ref_images: list of tensors, 每个tensor shape: [3, H, W]
        video_paths: list of str, 对应的视频路径
        first_frame_faces: tensor, dataset提供的crop人脸 [b, 3, face_H, face_W]
        crop_size: int, 裁剪后的人脸尺寸
        device: torch设备
        torch_dtype: torch数据类型
        global_rank: 全局rank，用于控制日志输出
        current_global_step: 当前训练步数
        
    Returns:
        list of tensors: 每个tensor shape: [3, crop_size, crop_size]
    """
    square_faces = []
    
    for i, (ref_image, video_path) in enumerate(zip(ref_images, video_paths)):
        try:
            # 将tensor转换为numpy数组用于人脸检测
            # ref_image shape: [3, H, W], 需要转换为 [H, W, 3]
            # 确保转换为float32类型，避免BFloat16问题
            ref_image_np = ref_image.permute(1, 2, 0).cpu().float().numpy()
            
            # 处理值域：如果值域是[-1,1]，则转换到[0,1]；如果已经是[0,1]，则直接使用
            if ref_image_np.min() < 0:
                # 从[-1,1]转换到[0,1]
                ref_image_np = (ref_image_np + 1) / 2
            ref_image_np = np.clip(ref_image_np, 0, 1)
            
            # 转换为uint8格式
            ref_image_np = (ref_image_np * 255).astype(np.uint8)
            
            # # 添加调试信息
            # if global_rank == 0 and current_global_step % 100 == 0:
            #     print(f"[Face Parser] 原始ref_image值域: [{ref_image.min():.4f}, {ref_image.max():.4f}]")
            #     print(f"[Face Parser] 转换后numpy数组值域: [{ref_image_np.min()}, {ref_image_np.max()}]")
            #     print(f"[Face Parser] 转换后numpy数组形状: {ref_image_np.shape}")
            
            # 使用人脸解析器检测人脸
            face_result = face_parser.infer_from_array(ref_image_np, n=1)  # 只取最大的人脸
            
            if face_result and 'bboxes' in face_result and len(face_result['bboxes']) > 0:
                # 获取第一个（最大的）人脸的bbox
                bbox = face_result['bboxes'][0]  # [x, y, width, height]
                x, y, w, h = bbox
                
                # 从原图中裁剪人脸区域
                face_crop = ref_image_np[int(y):int(y+h), int(x):int(x+w)]
                
                # # 添加调试信息
                # if global_rank == 0 and current_global_step % 100 == 0:
                #     print(f"[Face Parser] 成功提取人脸 {i}: bbox=[{x:.1f}, {y:.1f}, {w:.1f}, {h:.1f}]")
                #     print(f"[Face Parser] 裁剪后的人脸形状: {face_crop.shape}")
                #     print(f"[Face Parser] 裁剪后的人脸值域: [{face_crop.min()}, {face_crop.max()}]")
                
                # 将numpy数组转换回tensor并调整尺寸
                face_tensor = torch.from_numpy(face_crop).permute(2, 0, 1).float() / 255.0
                face_tensor = resize(face_tensor, size=(crop_size, crop_size), interpolation=InterpolationMode.BILINEAR)
                
                # # 添加更多调试信息
                # if global_rank == 0 and current_global_step % 100 == 0:
                #     print(f"[Face Parser] 调整尺寸后的tensor形状: {face_tensor.shape}")
                #     print(f"[Face Parser] 调整尺寸后的tensor值域: [{face_tensor.min():.4f}, {face_tensor.max():.4f}]")
                
                square_faces.append(face_tensor.to(device, dtype=torch_dtype))
                    
            else:
                # 如果没有检测到人脸，回退到使用dataset提供的crop人脸
                dataset_face = first_frame_faces[i]  # [3, face_H, face_W]
                
                # 添加调试信息
                if global_rank == 0 :
                    print(f"[Face Parser] 未检测到人脸 {i}，回退到dataset提供的crop人脸")
                    print(f"[Face Parser] Dataset人脸形状: {dataset_face.shape}")
                    print(f"[Face Parser] Dataset人脸值域: [{dataset_face.min():.4f}, {dataset_face.max():.4f}]")
                
                face_tensor = resize(dataset_face, size=(crop_size, crop_size), interpolation=InterpolationMode.BILINEAR)
                
                # 添加更多调试信息
                if global_rank == 0:
                    print(f"[Face Parser] 调整尺寸后的dataset人脸形状: {face_tensor.shape}")
                    print(f"[Face Parser] 调整尺寸后的dataset人脸值域: [{face_tensor.min():.4f}, {face_tensor.max():.4f}]")
                
                square_faces.append(face_tensor.to(device, dtype=torch_dtype))
                    
        except Exception as e:
            # 异常情况下回退到使用dataset提供的crop人脸
            dataset_face = first_frame_faces[i]  # [3, face_H, face_W]
            
            # 添加调试信息
            if global_rank == 0:
                print(f"[Face Parser] 处理图像 {i} 时出错: {str(e)}，回退到dataset提供的crop人脸")
                print(f"[Face Parser] 异常回退 - Dataset人脸形状: {dataset_face.shape}")
                print(f"[Face Parser] 异常回退 - Dataset人脸值域: [{dataset_face.min():.4f}, {dataset_face.max():.4f}]")
            
            face_tensor = resize(dataset_face, size=(crop_size, crop_size), interpolation=InterpolationMode.BILINEAR)
            
            # 添加更多调试信息
            if global_rank == 0:
                print(f"[Face Parser] 异常回退 - 调整尺寸后的tensor形状: {face_tensor.shape}")
                print(f"[Face Parser] 异常回退 - 调整尺寸后的tensor值域: [{face_tensor.min():.4f}, {face_tensor.max():.4f}]")
            
            square_faces.append(face_tensor.to(device, dtype=torch_dtype))
    
    return square_faces


def save_face_parser_debug_images(ref_cropped_list, ref_image, video_paths, first_frame_faces, current_step, crop_image_size, debug_dir="./logs/debug/ref", global_rank=0):
    """
    保存人脸解析器的调试图像到debug目录
    
    Args:
        ref_cropped_list: list of tensors, 人脸解析器生成的方形人脸
        ref_image: tensor, 原始参考图像
        video_paths: list of str, 视频路径
        first_frame_faces: tensor, dataset提供的crop人脸
        current_step: int, 当前训练步数
        crop_image_size: int, 裁剪图像尺寸
        debug_dir: str, 调试目录路径
    """
    try:
        # 只在rank 0时创建debug目录
        if global_rank == 0:
            os.makedirs(debug_dir, exist_ok=True)
            os.chmod(debug_dir, 0o777)
        
        # 保存每个batch的调试图像
        for batch_idx, (ref_cropped, video_path) in enumerate(zip(ref_cropped_list, video_paths)):
            # 只在rank 0时创建子目录
            if global_rank == 0:
                batch_dir = os.path.join(debug_dir, f"step_{current_step}_batch_{batch_idx}")
                os.makedirs(batch_dir, exist_ok=True)
                os.chmod(batch_dir, 0o777)
            
            # 保存原始参考图像
            if ref_image is not None:
                ref_img_single = ref_image[batch_idx]  # [3, H, W]
                # 处理值域：如果值域是[-1,1]，则转换到[0,1]；如果已经是[0,1]，则直接使用
                if ref_img_single.min() < 0:
                    # 从[-1,1]转换到[0,1]
                    ref_img_single = (ref_img_single + 1) / 2
                ref_img_single = torch.clamp(ref_img_single, 0, 1)
                ref_img_path = os.path.join(batch_dir, "original_ref_image.png")
                save_image(ref_img_single, ref_img_path)
            
            # 保存dataset提供的原始crop人脸
            if first_frame_faces is not None:
                dataset_face_single = first_frame_faces[batch_idx]  # [3, face_H, face_W]
                # 处理值域：如果值域是[-1,1]，则转换到[0,1]；如果已经是[0,1]，则直接使用
                if dataset_face_single.min() < 0:
                    # 从[-1,1]转换到[0,1]
                    dataset_face_single = (dataset_face_single + 1) / 2
                dataset_face_single = torch.clamp(dataset_face_single, 0, 1)
                dataset_face_path = os.path.join(batch_dir, "dataset_crop_face.png")
                save_image(dataset_face_single, dataset_face_path)
            
            # 保存人脸解析器生成的方形人脸
            if ref_cropped is not None:
                ref_cropped_single = ref_cropped[batch_idx]  # [3, crop_size, crop_size]
                
                # 添加调试信息
                print(f"[Debug] Parsed face tensor shape: {ref_cropped_single.shape}")
                print(f"[Debug] Parsed face tensor dtype: {ref_cropped_single.dtype}")
                print(f"[Debug] Parsed face tensor device: {ref_cropped_single.device}")
                print(f"[Debug] Parsed face tensor stats: min={ref_cropped_single.min():.4f}, max={ref_cropped_single.max():.4f}, mean={ref_cropped_single.mean():.4f}, std={ref_cropped_single.std():.4f}")
                
                # 检查是否有异常值
                if torch.isnan(ref_cropped_single).any():
                    print(f"[Debug] Warning: Parsed face contains NaN values!")
                if torch.isinf(ref_cropped_single).any():
                    print(f"[Debug] Warning: Parsed face contains Inf values!")
                
                # 处理值域：如果值域是[-1,1]，则转换到[0,1]；如果已经是[0,1]，则直接使用
                if ref_cropped_single.min() < 0:
                    # 从[-1,1]转换到[0,1]
                    ref_cropped_single = (ref_cropped_single + 1) / 2
                ref_cropped_single = torch.clamp(ref_cropped_single, 0, 1)
                
                # 保存处理后的图像
                face_path = os.path.join(batch_dir, "parsed_square_face.png")
                save_image(ref_cropped_single, face_path)
                
                # 额外保存一个未处理的版本用于对比
                raw_face_path = os.path.join(batch_dir, "parsed_square_face_raw.png")
                save_image(ref_cropped[batch_idx], raw_face_path)
            
            # 保存视频路径信息
            info_path = os.path.join(batch_dir, "info.txt")
            with open(info_path, 'w') as f:
                f.write(f"Video Path: {video_path}\n")
                f.write(f"Step: {current_step}\n")
                f.write(f"Batch Index: {batch_idx}\n")
                f.write(f"Face Size: {crop_image_size}x{crop_image_size}\n")
                if first_frame_faces is not None:
                    f.write(f"Dataset Face Size: {first_frame_faces[batch_idx].shape[1]}x{first_frame_faces[batch_idx].shape[2]}\n")
                
                # 添加tensor值域信息
                if ref_image is not None:
                    ref_img_single = ref_image[batch_idx]
                    f.write(f"Original Ref Image Range: [{ref_img_single.min():.4f}, {ref_img_single.max():.4f}]\n")
                
                if first_frame_faces is not None:
                    dataset_face_single = first_frame_faces[batch_idx]
                    f.write(f"Dataset Face Range: [{dataset_face_single.min():.4f}, {dataset_face_single.max():.4f}]\n")
                
                if ref_cropped is not None:
                    ref_cropped_single = ref_cropped[batch_idx]
                    f.write(f"Parsed Face Range: [{ref_cropped_single.min():.4f}, {ref_cropped_single.max():.4f}]\n")
            
            print(f"[Debug] 保存人脸解析器调试图像到: {batch_dir}")
            
    except Exception as e:
        print(f"[Debug] 保存人脸解析器调试图像时出错: {str(e)}")