LiveAvatar/liveavatar/util.py

from regex import F
from torch.distributed.fsdp.wrap import transformer_auto_wrap_policy
from torch.distributed.fsdp.wrap import size_based_auto_wrap_policy
from torch.distributed.fsdp.wrap import ModuleWrapPolicy
from torch.distributed.fsdp import (
    FullStateDictConfig,
    StateDictType,
    MixedPrecision,
    ShardingStrategy,
    FullyShardedDataParallel as FSDP
)
from torchvision.utils import make_grid
from datetime import timedelta, datetime
import torch.distributed as dist
from omegaconf import OmegaConf
from functools import partial
import numpy as np
import random
import torch
import os


def launch_distributed_job(backend: str = "nccl"):
    rank = int(os.environ["RANK"])
    local_rank = int(os.environ["LOCAL_RANK"])
    world_size = int(os.environ["WORLD_SIZE"])
    host = os.environ["MASTER_ADDR"]
    port = int(os.environ["MASTER_PORT"])

    if ":" in host:  # IPv6
        init_method = f"tcp://[{host}]:{port}"
    else:  # IPv4
        init_method = f"tcp://{host}:{port}"
    dist.init_process_group(rank=rank, world_size=world_size, backend=backend,
                            init_method=init_method, timeout=timedelta(minutes=30))
    torch.cuda.set_device(local_rank)


def set_seed(seed: int, deterministic: bool = False):
    """
    Helper function for reproducible behavior to set the seed in `random`, `numpy`, `torch`.

    Args:
        seed (`int`):
            The seed to set.
        deterministic (`bool`, *optional*, defaults to `False`):
            Whether to use deterministic algorithms where available. Can slow down training.
    """
    random.seed(seed)
    np.random.seed(seed)
    torch.manual_seed(seed)
    torch.cuda.manual_seed_all(seed)

    if deterministic:
        torch.use_deterministic_algorithms(True)


def init_logging_folder(args):
    import wandb
    date = str(datetime.now()).replace(" ", "-").replace(":", "-")
    output_path = os.path.join(
        args.output_path,
        f"{date}_seed{args.seed}"
    )
    os.makedirs(output_path, exist_ok=False)

    os.makedirs(args.output_path, exist_ok=True)
    wandb.login(host=args.wandb_host, key=args.wandb_key)
    run = wandb.init(config=OmegaConf.to_container(args, resolve=True), dir=args.output_path, **
                     {"mode": "online", "entity": args.wandb_entity, "project": args.wandb_project})
    wandb.run.log_code(".")
    wandb.run.name = args.wandb_name
    print(f"run dir: {run.dir}")
    wandb_folder = run.dir
    os.makedirs(wandb_folder, exist_ok=True)

    return output_path, wandb_folder


def fsdp_wrap(module, sharding_strategy="full", mixed_precision=False, wrap_strategy="size", min_num_params=int(5e7), transformer_module=None):
    if mixed_precision:
        mixed_precision_policy = MixedPrecision(
            param_dtype=torch.bfloat16,
            reduce_dtype=torch.float32,
            buffer_dtype=torch.float32,
            cast_forward_inputs=False
        )
    else:
        mixed_precision_policy = None

    if wrap_strategy == "transformer":
        auto_wrap_policy = partial(
            transformer_auto_wrap_policy,
            transformer_layer_cls=transformer_module
        )
    elif wrap_strategy == "size":
        auto_wrap_policy = partial(
            size_based_auto_wrap_policy,
            min_num_params=min_num_params
        )
    else:
        raise ValueError(f"Invalid wrap strategy: {wrap_strategy}")

    os.environ["NCCL_CROSS_NIC"] = "1"

    sharding_strategy = {
        "full": ShardingStrategy.FULL_SHARD,
        "hybrid_full": ShardingStrategy.HYBRID_SHARD,
        "hybrid_zero2": ShardingStrategy._HYBRID_SHARD_ZERO2,
        "no_shard": ShardingStrategy.NO_SHARD,
    }[sharding_strategy]

    module = FSDP(
        module,
        auto_wrap_policy=auto_wrap_policy,
        sharding_strategy=sharding_strategy,
        mixed_precision=mixed_precision_policy,
        device_id=torch.cuda.current_device(),
        limit_all_gathers=True,
        sync_module_states=False  # Load ckpt on rank 0 and sync to other ranks
    )
    return module


def cycle(dl):
    while True:
        for data in dl:
            yield data


def fsdp_state_dict(model):
    fsdp_fullstate_save_policy = FullStateDictConfig(
        offload_to_cpu=True, rank0_only=True
    )
    with FSDP.state_dict_type(
        model, StateDictType.FULL_STATE_DICT, fsdp_fullstate_save_policy
    ):
        checkpoint = model.state_dict()

    return checkpoint


def barrier():
    if dist.is_initialized():
        dist.barrier()


def prepare_for_saving(tensor, fps=16, caption=None):
    import wandb
    # Convert range [-1, 1] to [0, 1]
    tensor = (tensor * 0.5 + 0.5).clamp(0, 1).detach()

    if tensor.ndim == 4:
        # Assuming it's an image and has shape [batch_size, 3, height, width]
        tensor = make_grid(tensor, 4, padding=0, normalize=False)
        return wandb.Image((tensor * 255).cpu().numpy().astype(np.uint8), caption=caption)
    elif tensor.ndim == 5:
        # Assuming it's a video and has shape [batch_size, num_frames, 3, height, width]
        return wandb.Video((tensor * 255).cpu().numpy().astype(np.uint8), fps=fps, format="webm", caption=caption)
    else:
        raise ValueError("Unsupported tensor shape for saving. Expected 4D (image) or 5D (video) tensor.")

def validate_tensor(tensor, expected_shape, expected_dtype, expected_device):
    """验证单个tensor的形状、dtype和device"""
    try:
        if not isinstance(tensor, torch.Tensor):
            raise Exception("tensor is not a torch.Tensor")
        if tensor.shape != expected_shape:
            raise Exception("tensor shape is not expected")
        if tensor.dtype != expected_dtype:
            raise Exception("tensor dtype is not expected")
        if tensor.device != expected_device:
            raise Exception("tensor device is not expected")
    except Exception as e:
        # import pdb;pdb.set_trace()
        print(f"error: {e}")
        return False
    return True


def check_input(noise_input,timestep,arg_c):
    """
    验证输入参数的tensor是否符合预期规格
    
    Args:
        arg_c: 包含各种tensor的字典
        
    Returns:
        bool: 所有验证通过返回True，否则返回False
    """
    try:
        bs = noise_input.__len__()      # batch_size
        noise_input = noise_input[0]      # [c,f,h,w]
        f = noise_input.shape[1]       # frames  
        h = noise_input.shape[2]       # height
        w = noise_input.shape[3]       # width
        motion_frames_rgb = arg_c["motion_frames"][0]
        motion_frames_latent = arg_c["motion_frames"][1]
        
        expected_dtype = noise_input.dtype
        expected_device = noise_input.device
        
        validation_configs = [
            ("cond_states", (bs, 16, f, h, w), True),
            ("ref_latents", (bs, 16,1, h, w), True),
            ("motion_latents", (bs, 16, motion_frames_latent, h, w), False),
            ("audio_input", (bs, 25, 1024, 4*f), True),#数据集里 1的部分在 motion_frames，后面都是 4n
            ("context", (bs, 512, 4096), True),
        ]
        if validate_tensor(timestep, (bs, f), expected_dtype, expected_device) is False:
            print(f"input timestep, shape {timestep.shape}, dtype {timestep.dtype}, device {timestep.device} is not valid")
            print(f"expected shape {bs, f}, expected dtype {expected_dtype}, expected device {expected_device}")
            return False

        for key, expected_shape, is_required in validation_configs:
            if key in arg_c:
                if not validate_tensor(arg_c[key], expected_shape, expected_dtype, expected_device):
                    print(f"input {key}, shape {arg_c[key].shape}, dtype {arg_c[key].dtype}, device {arg_c[key].device} is not valid")
                    print(f"expected shape {expected_shape}, expected dtype {expected_dtype}, expected device {expected_device}")
                    return False
            elif is_required:
                print(f"input {key} is required, but not found")
                return False
                print(f"input {key} is required, but not found")
    except Exception as e:
        print(f"error: {e}")
        return False
 
    return True