wan/utils/utils.py

import argparse
import binascii
import os
import os.path as osp
import imageio
import torch
import torchvision
__all__ = ['cache_video', 'cache_image', 'str2bool']

def rand_name(length=8, suffix=''):
    name = binascii.b2a_hex(os.urandom(length)).decode('utf-8')
    if suffix:
        if not suffix.startswith('.'):
            suffix = '.' + suffix
        name += suffix
    return name

def cache_video(tensor, save_file=None, fps=30, suffix='.mp4', nrow=8, normalize=True, value_range=(-1, 1), retry=5):
    cache_file = osp.join('/tmp', rand_name(suffix=suffix)) if save_file is None else save_file
    error = None
    for _ in range(retry):
        try:
            tensor = tensor.clamp(min(value_range), max(value_range))
            tensor = torch.stack([torchvision.utils.make_grid(u, nrow=nrow, normalize=normalize, value_range=value_range) for u in tensor.unbind(2)], dim=1).permute(1, 2, 3, 0)
            tensor = (tensor * 255).type(torch.uint8).cpu()
            writer = imageio.get_writer(cache_file, fps=fps, codec='libx264', quality=8)
            for frame in tensor.numpy():
                writer.append_data(frame)
            writer.close()
            return cache_file
        except Exception as e:
            error = e
            continue
    else:
        print(f'cache_video failed, error: {error}', flush=True)
        return None

def cache_image(tensor, save_file, nrow=8, normalize=True, value_range=(-1, 1), retry=5):
    suffix = osp.splitext(save_file)[1]
    if suffix.lower() not in ['.jpg', '.jpeg', '.png', '.tiff', '.gif', '.webp']:
        suffix = '.png'
    error = None
    for _ in range(retry):
        try:
            tensor = tensor.clamp(min(value_range), max(value_range))
            torchvision.utils.save_image(tensor, save_file, nrow=nrow, normalize=normalize, value_range=value_range)
            return save_file
        except Exception as e:
            error = e
            continue

def str2bool(v):
    if isinstance(v, bool):
        return v
    v_lower = v.lower()
    if v_lower in ('yes', 'true', 't', 'y', '1'):
        return True
    elif v_lower in ('no', 'false', 'f', 'n', '0'):
        return False
    else:
        raise argparse.ArgumentTypeError('Boolean value expected (True/False)')