957 lines
31 KiB
Python
957 lines
31 KiB
Python
# These codes are copied from modelscope revision c58451baead80d83281f063d12fb377fad415257
|
|
# Copyright (c) Alibaba, Inc. and its affiliates.
|
|
|
|
import base64
|
|
import binascii
|
|
import copy
|
|
import glob
|
|
import gzip
|
|
import hashlib
|
|
import logging
|
|
import math
|
|
import os
|
|
import os.path as osp
|
|
import pickle
|
|
import random
|
|
import sys
|
|
import time
|
|
import urllib.request
|
|
import zipfile
|
|
from io import BytesIO
|
|
from multiprocessing.pool import ThreadPool as Pool
|
|
|
|
import imageio
|
|
import json
|
|
import numpy as np
|
|
import oss2 as oss
|
|
import requests
|
|
import skvideo.io
|
|
import torch
|
|
import torch.nn.functional as F
|
|
import torchvision.utils as tvutils
|
|
from einops import rearrange
|
|
from PIL import Image
|
|
|
|
__all__ = [
|
|
'parse_oss_url', 'parse_bucket', 'read', 'read_image', 'read_gzip',
|
|
'ceil_divide', 'to_device', 'put_object', 'put_torch_object',
|
|
'put_object_from_file', 'get_object', 'get_object_to_file', 'rand_name',
|
|
'save_image', 'save_video', 'save_video_vs_conditions',
|
|
'save_video_multiple_conditions_with_data',
|
|
'save_video_multiple_conditions', 'download_video_to_file',
|
|
'save_video_grid_mp4', 'save_caps', 'ema', 'parallel', 'exists',
|
|
'download', 'unzip', 'load_state_dict', 'inverse_indices',
|
|
'detect_duplicates', 'read_tfs', 'md5', 'rope', 'format_state',
|
|
'breakup_grid', 'huggingface_tokenizer', 'huggingface_model'
|
|
]
|
|
|
|
TFS_CLIENT = None
|
|
|
|
|
|
def DOWNLOAD_TO_CACHE(oss_key,
|
|
file_or_dirname=None,
|
|
cache_dir=osp.join(
|
|
'/'.join(osp.abspath(__file__).split('/')[:-2]),
|
|
'model_weights')):
|
|
r"""Download OSS [file or folder] to the cache folder.
|
|
Only the 0th process on each node will run the downloading.
|
|
Barrier all processes until the downloading is completed.
|
|
"""
|
|
# source and target paths
|
|
base_path = osp.join(cache_dir, file_or_dirname or osp.basename(oss_key))
|
|
|
|
return base_path
|
|
|
|
|
|
def find_free_port():
|
|
"""https://stackoverflow.com/questions/1365265/on-localhost-how-do-i-pick-a-free-port-number"""
|
|
import socket
|
|
from contextlib import closing
|
|
with closing(socket.socket(socket.AF_INET, socket.SOCK_STREAM)) as s:
|
|
s.bind(('', 0))
|
|
s.setsockopt(socket.SOL_SOCKET, socket.SO_REUSEADDR, 1)
|
|
return str(s.getsockname()[1])
|
|
|
|
|
|
def setup_seed(seed):
|
|
torch.manual_seed(seed)
|
|
torch.cuda.manual_seed_all(seed)
|
|
np.random.seed(seed)
|
|
random.seed(seed)
|
|
torch.backends.cudnn.deterministic = True
|
|
|
|
|
|
def parse_oss_url(path):
|
|
if path.startswith('oss://'):
|
|
path = path[len('oss://'):]
|
|
|
|
# configs
|
|
configs = {
|
|
'endpoint': os.getenv('OSS_ENDPOINT', None),
|
|
'accessKeyID': os.getenv('OSS_ACCESS_KEY_ID', None),
|
|
'accessKeySecret': os.getenv('OSS_ACCESS_KEY_SECRET', None),
|
|
'securityToken': os.getenv('OSS_SECURITY_TOKEN', None)
|
|
}
|
|
bucket, path = path.split('/', maxsplit=1)
|
|
if '?' in bucket:
|
|
bucket, config = bucket.split('?', maxsplit=1)
|
|
for pair in config.split('&'):
|
|
k, v = pair.split('=', maxsplit=1)
|
|
configs[k] = v
|
|
|
|
# session
|
|
session = parse_oss_url._sessions.setdefault(f'{bucket}@{os.getpid()}',
|
|
oss.Session())
|
|
|
|
# bucket
|
|
bucket = oss.Bucket(
|
|
auth=oss.Auth(configs['accessKeyID'], configs['accessKeySecret']),
|
|
endpoint=configs['endpoint'],
|
|
bucket_name=bucket,
|
|
session=session)
|
|
return bucket, path
|
|
|
|
|
|
parse_oss_url._sessions = {}
|
|
|
|
|
|
def parse_bucket(url):
|
|
return parse_oss_url(osp.join(url, '_placeholder'))[0]
|
|
|
|
|
|
def read(filename, mode='r', retry=5):
|
|
assert mode in ['r', 'rb']
|
|
exception = None
|
|
for _ in range(retry):
|
|
try:
|
|
if filename.startswith('oss://'):
|
|
bucket, path = parse_oss_url(filename)
|
|
content = bucket.get_object(path).read()
|
|
if mode == 'r':
|
|
content = content.decode('utf-8')
|
|
elif filename.startswith('http'):
|
|
content = requests.get(filename).content
|
|
if mode == 'r':
|
|
content = content.decode('utf-8')
|
|
else:
|
|
with open(filename, mode=mode) as f:
|
|
content = f.read()
|
|
return content
|
|
except Exception as e:
|
|
exception = e
|
|
continue
|
|
else:
|
|
raise exception
|
|
|
|
|
|
def read_image(filename, retry=5):
|
|
exception = None
|
|
for _ in range(retry):
|
|
try:
|
|
return Image.open(BytesIO(read(filename, mode='rb', retry=retry)))
|
|
except Exception as e:
|
|
exception = e
|
|
continue
|
|
else:
|
|
raise exception
|
|
|
|
|
|
def download_video_to_file(filename, local_file, retry=5):
|
|
exception = None
|
|
for _ in range(retry):
|
|
try:
|
|
bucket, path = parse_oss_url(filename)
|
|
bucket.get_object_to_file(path, local_file)
|
|
break
|
|
except Exception as e:
|
|
exception = e
|
|
continue
|
|
else:
|
|
raise exception
|
|
|
|
|
|
def read_gzip(filename, retry=5):
|
|
exception = None
|
|
for _ in range(retry):
|
|
try:
|
|
remove = False
|
|
if filename.startswith('oss://'):
|
|
bucket, path = parse_oss_url(filename)
|
|
filename = rand_name(suffix=osp.splitext(filename)[1])
|
|
bucket.get_object_to_file(path, filename)
|
|
remove = True
|
|
with gzip.open(filename) as f:
|
|
content = f.read()
|
|
if remove:
|
|
os.remove(filename)
|
|
return content
|
|
except Exception as e:
|
|
exception = e
|
|
continue
|
|
else:
|
|
raise exception
|
|
|
|
|
|
def ceil_divide(a, b):
|
|
return int(math.ceil(a / b))
|
|
|
|
|
|
def to_device(batch, device, non_blocking=False):
|
|
if isinstance(batch, (list, tuple)):
|
|
return type(batch)([to_device(u, device, non_blocking) for u in batch])
|
|
elif isinstance(batch, dict):
|
|
return type(batch)([(k, to_device(v, device, non_blocking))
|
|
for k, v in batch.items()])
|
|
elif isinstance(batch, torch.Tensor) and batch.device != device:
|
|
batch = batch.to(device, non_blocking=non_blocking)
|
|
return batch
|
|
|
|
|
|
def put_object(bucket, oss_key, data, retry=5):
|
|
exception = None
|
|
for _ in range(retry):
|
|
try:
|
|
return bucket.put_object(oss_key, data)
|
|
except Exception as e:
|
|
exception = e
|
|
continue
|
|
else:
|
|
print(
|
|
f'put_object to {oss_key} failed with error: {exception}',
|
|
flush=True)
|
|
|
|
|
|
def put_torch_object(bucket, oss_key, data, retry=5):
|
|
exception = None
|
|
for _ in range(retry):
|
|
try:
|
|
buffer = BytesIO()
|
|
torch.save(data, buffer)
|
|
return bucket.put_object(oss_key, buffer.getvalue())
|
|
except Exception as e:
|
|
exception = e
|
|
continue
|
|
else:
|
|
print(
|
|
f'put_torch_object to {oss_key} failed with error: {exception}',
|
|
flush=True)
|
|
|
|
|
|
def put_object_from_file(bucket, oss_key, filename, retry=5):
|
|
exception = None
|
|
for _ in range(retry):
|
|
try:
|
|
return bucket.put_object_from_file(oss_key, filename)
|
|
except Exception as e:
|
|
exception = e
|
|
continue
|
|
else:
|
|
print(
|
|
f'put_object_from_file to {oss_key} failed with error: {exception}',
|
|
flush=True)
|
|
|
|
|
|
def get_object(bucket, oss_key, retry=5):
|
|
exception = None
|
|
for _ in range(retry):
|
|
try:
|
|
return bucket.get_object(oss_key).read()
|
|
except Exception as e:
|
|
exception = e
|
|
continue
|
|
else:
|
|
print(
|
|
f'get_object from {oss_key} failed with error: {exception}',
|
|
flush=True)
|
|
|
|
|
|
def get_object_to_file(bucket, oss_key, filename, retry=5):
|
|
exception = None
|
|
for _ in range(retry):
|
|
try:
|
|
return bucket.get_object_to_file(oss_key, filename)
|
|
except Exception as e:
|
|
exception = e
|
|
continue
|
|
else:
|
|
print(
|
|
f'get_object_to_file from {oss_key} failed with error: {exception}',
|
|
flush=True)
|
|
|
|
|
|
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
|
|
|
|
|
|
@torch.no_grad()
|
|
def save_image(bucket,
|
|
oss_key,
|
|
tensor,
|
|
nrow=8,
|
|
normalize=True,
|
|
range=(-1, 1),
|
|
retry=5):
|
|
filename = rand_name(suffix='.jpg')
|
|
for _ in [None] * retry:
|
|
try:
|
|
tvutils.save_image(
|
|
tensor, filename, nrow=nrow, normalize=normalize, range=range)
|
|
bucket.put_object_from_file(oss_key, filename)
|
|
exception = None
|
|
break
|
|
except Exception as e:
|
|
exception = e
|
|
continue
|
|
|
|
# remove temporary file
|
|
if osp.exists(filename):
|
|
os.remove(filename)
|
|
if exception is not None:
|
|
print(
|
|
'save image to {} failed, error: {}'.format(oss_key, exception),
|
|
flush=True)
|
|
|
|
|
|
@torch.no_grad()
|
|
def video_tensor_to_gif(tensor, path, duration=120, loop=0, optimize=True):
|
|
tensor = tensor.permute(1, 2, 3, 0)
|
|
images = tensor.unbind(dim=0)
|
|
images = [(image.numpy() * 255).astype('uint8') for image in images]
|
|
imageio.mimwrite(path, images, fps=8)
|
|
return images
|
|
|
|
|
|
@torch.no_grad()
|
|
def save_video(bucket,
|
|
oss_key,
|
|
tensor,
|
|
mean=[0.5, 0.5, 0.5],
|
|
std=[0.5, 0.5, 0.5],
|
|
nrow=8,
|
|
retry=5):
|
|
mean = torch.tensor(mean, device=tensor.device).view(1, -1, 1, 1, 1)
|
|
std = torch.tensor(std, device=tensor.device).view(1, -1, 1, 1, 1)
|
|
tensor = tensor.mul_(std).add_(mean)
|
|
tensor.clamp_(0, 1)
|
|
|
|
filename = rand_name(suffix='.gif')
|
|
for _ in [None] * retry:
|
|
try:
|
|
one_gif = rearrange(
|
|
tensor, '(i j) c f h w -> c f (i h) (j w)', i=nrow)
|
|
video_tensor_to_gif(one_gif, filename)
|
|
bucket.put_object_from_file(oss_key, filename)
|
|
exception = None
|
|
break
|
|
except Exception as e:
|
|
exception = e
|
|
continue
|
|
|
|
# remove temporary file
|
|
if osp.exists(filename):
|
|
os.remove(filename)
|
|
if exception is not None:
|
|
print(
|
|
'save video to {} failed, error: {}'.format(oss_key, exception),
|
|
flush=True)
|
|
|
|
|
|
@torch.no_grad()
|
|
def save_video_multiple_conditions(oss_key,
|
|
video_tensor,
|
|
model_kwargs,
|
|
source_imgs,
|
|
palette,
|
|
mean=[0.5, 0.5, 0.5],
|
|
std=[0.5, 0.5, 0.5],
|
|
nrow=8,
|
|
retry=5,
|
|
save_origin_video=True,
|
|
bucket=None):
|
|
mean = torch.tensor(mean, device=video_tensor.device).view(1, -1, 1, 1, 1)
|
|
std = torch.tensor(std, device=video_tensor.device).view(1, -1, 1, 1, 1)
|
|
video_tensor = video_tensor.mul_(std).add_(mean)
|
|
try:
|
|
video_tensor.clamp_(0, 1)
|
|
except Exception as e:
|
|
video_tensor = video_tensor.float().clamp_(0, 1)
|
|
print(e)
|
|
video_tensor = video_tensor.cpu()
|
|
|
|
b, c, n, h, w = video_tensor.shape
|
|
source_imgs = F.adaptive_avg_pool3d(source_imgs, (n, h, w))
|
|
source_imgs = source_imgs.cpu()
|
|
|
|
model_kwargs_channel3 = {}
|
|
for key, conditions in model_kwargs[0].items():
|
|
if conditions.shape[-1] == 1024:
|
|
# Skip for style embeding
|
|
continue
|
|
if len(conditions.shape) == 3:
|
|
conditions_np = conditions.cpu().numpy()
|
|
conditions = []
|
|
for i in conditions_np:
|
|
vis_i = []
|
|
for j in i:
|
|
vis_i.append(
|
|
palette.get_palette_image(
|
|
j, percentile=90, width=256, height=256))
|
|
conditions.append(np.stack(vis_i))
|
|
conditions = torch.from_numpy(np.stack(conditions))
|
|
conditions = rearrange(conditions, 'b n h w c -> b c n h w')
|
|
else:
|
|
if conditions.size(1) == 1:
|
|
conditions = torch.cat([conditions, conditions, conditions],
|
|
dim=1)
|
|
conditions = F.adaptive_avg_pool3d(conditions, (n, h, w))
|
|
if conditions.size(1) == 2:
|
|
conditions = torch.cat([conditions, conditions[:, :1, ]],
|
|
dim=1)
|
|
conditions = F.adaptive_avg_pool3d(conditions, (n, h, w))
|
|
elif conditions.size(1) == 3:
|
|
conditions = F.adaptive_avg_pool3d(conditions, (n, h, w))
|
|
elif conditions.size(1) == 4:
|
|
color = ((conditions[:, 0:3] + 1.) / 2.)
|
|
alpha = conditions[:, 3:4]
|
|
conditions = color * alpha + 1.0 * (1.0 - alpha)
|
|
conditions = F.adaptive_avg_pool3d(conditions, (n, h, w))
|
|
model_kwargs_channel3[key] = conditions.cpu(
|
|
) if conditions.is_cuda else conditions
|
|
|
|
filename = oss_key
|
|
for _ in [None] * retry:
|
|
try:
|
|
vid_gif = rearrange(
|
|
video_tensor, '(i j) c f h w -> c f (i h) (j w)', i=nrow)
|
|
cons_list = [
|
|
rearrange(con, '(i j) c f h w -> c f (i h) (j w)', i=nrow)
|
|
for _, con in model_kwargs_channel3.items()
|
|
]
|
|
source_imgs = rearrange(
|
|
source_imgs, '(i j) c f h w -> c f (i h) (j w)', i=nrow)
|
|
|
|
if save_origin_video:
|
|
vid_gif = torch.cat(
|
|
[
|
|
source_imgs,
|
|
] + cons_list + [
|
|
vid_gif,
|
|
], dim=3)
|
|
else:
|
|
vid_gif = torch.cat(
|
|
cons_list + [
|
|
vid_gif,
|
|
], dim=3)
|
|
|
|
video_tensor_to_gif(vid_gif, filename)
|
|
exception = None
|
|
break
|
|
except Exception as e:
|
|
exception = e
|
|
continue
|
|
if exception is not None:
|
|
logging.info('save video to {} failed, error: {}'.format(
|
|
oss_key, exception))
|
|
|
|
|
|
@torch.no_grad()
|
|
def save_video_multiple_conditions_with_data(bucket,
|
|
video_save_key,
|
|
gt_video_save_key,
|
|
vis_oss_key,
|
|
video_tensor,
|
|
model_kwargs,
|
|
source_imgs,
|
|
palette,
|
|
mean=[0.5, 0.5, 0.5],
|
|
std=[0.5, 0.5, 0.5],
|
|
nrow=8,
|
|
retry=5):
|
|
mean = torch.tensor(mean, device=video_tensor.device).view(1, -1, 1, 1, 1)
|
|
std = torch.tensor(std, device=video_tensor.device).view(1, -1, 1, 1, 1)
|
|
video_tensor = video_tensor.mul_(std).add_(mean)
|
|
video_tensor.clamp_(0, 1)
|
|
|
|
b, c, n, h, w = video_tensor.shape
|
|
source_imgs = F.adaptive_avg_pool3d(source_imgs, (n, h, w))
|
|
source_imgs = source_imgs.cpu()
|
|
|
|
model_kwargs_channel3 = {}
|
|
for key, conditions in model_kwargs[0].items():
|
|
if len(conditions.shape) == 3:
|
|
conditions_np = conditions.cpu().numpy()
|
|
conditions = []
|
|
for i in conditions_np:
|
|
vis_i = []
|
|
for j in i:
|
|
vis_i.append(
|
|
palette.get_palette_image(
|
|
j, percentile=90, width=256, height=256))
|
|
conditions.append(np.stack(vis_i))
|
|
conditions = torch.from_numpy(np.stack(conditions))
|
|
conditions = rearrange(conditions, 'b n h w c -> b c n h w')
|
|
else:
|
|
if conditions.size(1) == 1:
|
|
conditions = torch.cat([conditions, conditions, conditions],
|
|
dim=1)
|
|
conditions = F.adaptive_avg_pool3d(conditions, (n, h, w))
|
|
if conditions.size(1) == 2:
|
|
conditions = torch.cat([conditions, conditions[:, :1, ]],
|
|
dim=1)
|
|
conditions = F.adaptive_avg_pool3d(conditions, (n, h, w))
|
|
elif conditions.size(1) == 3:
|
|
conditions = F.adaptive_avg_pool3d(conditions, (n, h, w))
|
|
elif conditions.size(1) == 4:
|
|
color = ((conditions[:, 0:3] + 1.) / 2.)
|
|
alpha = conditions[:, 3:4]
|
|
conditions = color * alpha + 1.0 * (1.0 - alpha)
|
|
conditions = F.adaptive_avg_pool3d(conditions, (n, h, w))
|
|
model_kwargs_channel3[key] = conditions.cpu(
|
|
) if conditions.is_cuda else conditions
|
|
|
|
copy_video_tensor = video_tensor.clone()
|
|
copy_source_imgs = source_imgs.clone()
|
|
|
|
filename = rand_name(suffix='.gif')
|
|
for _ in [None] * retry:
|
|
try:
|
|
vid_gif = rearrange(
|
|
video_tensor, '(i j) c f h w -> c f (i h) (j w)', i=nrow)
|
|
cons_list = [
|
|
rearrange(con, '(i j) c f h w -> c f (i h) (j w)', j=nrow)
|
|
for _, con in model_kwargs_channel3.items()
|
|
]
|
|
source_imgs = rearrange(
|
|
source_imgs, '(i j) c f h w -> c f (i h) (j w)', i=nrow)
|
|
vid_gif = torch.cat(
|
|
[
|
|
source_imgs,
|
|
] + cons_list + [
|
|
vid_gif,
|
|
], dim=3)
|
|
|
|
video_tensor_to_gif(vid_gif, filename)
|
|
bucket.put_object_from_file(vis_oss_key, filename)
|
|
exception = None
|
|
break
|
|
except Exception as e:
|
|
exception = e
|
|
continue
|
|
|
|
# remove temporary file
|
|
if osp.exists(filename):
|
|
os.remove(filename)
|
|
|
|
filename_pred = rand_name(suffix='.pkl')
|
|
for _ in [None] * retry:
|
|
try:
|
|
copy_video_np = (copy_video_tensor.numpy() * 255).astype('uint8')
|
|
pickle.dump(copy_video_np, open(filename_pred, 'wb'))
|
|
bucket.put_object_from_file(video_save_key, filename_pred)
|
|
break
|
|
except Exception as e:
|
|
print('error! ', video_save_key, e)
|
|
continue
|
|
|
|
# remove temporary file
|
|
if osp.exists(filename_pred):
|
|
os.remove(filename_pred)
|
|
|
|
filename_gt = rand_name(suffix='.pkl')
|
|
for _ in [None] * retry:
|
|
try:
|
|
copy_source_np = (copy_source_imgs.numpy() * 255).astype('uint8')
|
|
pickle.dump(copy_source_np, open(filename_gt, 'wb'))
|
|
bucket.put_object_from_file(gt_video_save_key, filename_gt)
|
|
break
|
|
except Exception as e:
|
|
print('error! ', gt_video_save_key, e)
|
|
continue
|
|
|
|
# remove temporary file
|
|
if osp.exists(filename_gt):
|
|
os.remove(filename_gt)
|
|
|
|
if exception is not None:
|
|
print(
|
|
'save video to {} failed, error: {}'.format(
|
|
vis_oss_key, exception),
|
|
flush=True)
|
|
|
|
|
|
@torch.no_grad()
|
|
def save_video_vs_conditions(bucket,
|
|
oss_key,
|
|
video_tensor,
|
|
conditions,
|
|
source_imgs,
|
|
mean=[0.5, 0.5, 0.5],
|
|
std=[0.5, 0.5, 0.5],
|
|
nrow=8,
|
|
retry=5):
|
|
mean = torch.tensor(mean, device=video_tensor.device).view(1, -1, 1, 1, 1)
|
|
std = torch.tensor(std, device=video_tensor.device).view(1, -1, 1, 1, 1)
|
|
video_tensor = video_tensor.mul_(std).add_(mean)
|
|
video_tensor.clamp_(0, 1)
|
|
|
|
b, c, n, h, w = video_tensor.shape
|
|
source_imgs = F.adaptive_avg_pool3d(source_imgs, (n, h, w))
|
|
source_imgs = source_imgs.cpu()
|
|
|
|
if conditions.size(1) == 1:
|
|
conditions = torch.cat([conditions, conditions, conditions], dim=1)
|
|
conditions = F.adaptive_avg_pool3d(conditions, (n, h, w))
|
|
|
|
filename = rand_name(suffix='.gif')
|
|
for _ in [None] * retry:
|
|
try:
|
|
vid_gif = rearrange(
|
|
video_tensor, '(i j) c f h w -> c f (i h) (j w)', i=nrow)
|
|
con_gif = rearrange(
|
|
conditions, '(i j) c f h w -> c f (i h) (j w)', i=nrow)
|
|
source_imgs = rearrange(
|
|
source_imgs, '(i j) c f h w -> c f (i h) (j w)', i=nrow)
|
|
vid_gif = torch.cat([vid_gif, con_gif, source_imgs], dim=2)
|
|
|
|
video_tensor_to_gif(vid_gif, filename)
|
|
bucket.put_object_from_file(oss_key, filename)
|
|
exception = None
|
|
break
|
|
except Exception as e:
|
|
exception = e
|
|
continue
|
|
|
|
# remove temporary file
|
|
if osp.exists(filename):
|
|
os.remove(filename)
|
|
if exception is not None:
|
|
print(
|
|
'save video to {} failed, error: {}'.format(oss_key, exception),
|
|
flush=True)
|
|
|
|
|
|
@torch.no_grad()
|
|
def save_video_grid_mp4(bucket,
|
|
oss_key,
|
|
tensor,
|
|
mean=[0.5, 0.5, 0.5],
|
|
std=[0.5, 0.5, 0.5],
|
|
nrow=None,
|
|
fps=5,
|
|
retry=5):
|
|
mean = torch.tensor(mean, device=tensor.device).view(1, -1, 1, 1, 1)
|
|
std = torch.tensor(std, device=tensor.device).view(1, -1, 1, 1, 1)
|
|
tensor = tensor.mul_(std).add_(mean)
|
|
tensor.clamp_(0, 1)
|
|
b, c, t, h, w = tensor.shape
|
|
tensor = tensor.permute(0, 2, 3, 4, 1)
|
|
tensor = (tensor.cpu().numpy() * 255).astype('uint8')
|
|
|
|
filename = rand_name(suffix='.mp4')
|
|
for _ in [None] * retry:
|
|
try:
|
|
if nrow is None:
|
|
nrow = math.ceil(math.sqrt(b))
|
|
ncol = math.ceil(b / nrow)
|
|
padding = 1
|
|
video_grid = np.zeros((t, (padding + h) * nrow + padding,
|
|
(padding + w) * ncol + padding, c),
|
|
dtype='uint8')
|
|
for i in range(b):
|
|
r = i // ncol
|
|
c_ = i % ncol
|
|
|
|
start_r = (padding + h) * r
|
|
start_c = (padding + w) * c_
|
|
video_grid[:, start_r:start_r + h,
|
|
start_c:start_c + w] = tensor[i]
|
|
skvideo.io.vwrite(filename, video_grid, inputdict={'-r': str(fps)})
|
|
|
|
bucket.put_object_from_file(oss_key, filename)
|
|
exception = None
|
|
break
|
|
except Exception as e:
|
|
exception = e
|
|
continue
|
|
|
|
# remove temporary file
|
|
if osp.exists(filename):
|
|
os.remove(filename)
|
|
if exception is not None:
|
|
print(
|
|
'save video to {} failed, error: {}'.format(oss_key, exception),
|
|
flush=True)
|
|
|
|
|
|
@torch.no_grad()
|
|
def save_text(bucket, oss_key, tensor, nrow=8, retry=5):
|
|
len = tensor.shape[0]
|
|
num_per_row = int(len / nrow)
|
|
assert (len == nrow * num_per_row)
|
|
texts = ''
|
|
for i in range(nrow):
|
|
for j in range(num_per_row):
|
|
text = dec_bytes2obj(tensor[i * num_per_row + j])
|
|
texts += text + '\n'
|
|
texts += '\n'
|
|
|
|
for _ in [None] * retry:
|
|
try:
|
|
bucket.put_object(oss_key, texts)
|
|
exception = None
|
|
break
|
|
except Exception as e:
|
|
exception = e
|
|
continue
|
|
if exception is not None:
|
|
print(
|
|
'save video to {} failed, error: {}'.format(oss_key, exception),
|
|
flush=True)
|
|
|
|
|
|
@torch.no_grad()
|
|
def save_caps(bucket, oss_key, caps, retry=5):
|
|
texts = ''
|
|
for cap in caps:
|
|
texts += cap
|
|
texts += '\n'
|
|
|
|
for _ in [None] * retry:
|
|
try:
|
|
bucket.put_object(oss_key, texts)
|
|
exception = None
|
|
break
|
|
except Exception as e:
|
|
exception = e
|
|
continue
|
|
if exception is not None:
|
|
print(
|
|
'save video to {} failed, error: {}'.format(oss_key, exception),
|
|
flush=True)
|
|
|
|
|
|
@torch.no_grad()
|
|
def ema(net_ema, net, beta, copy_buffer=False):
|
|
assert 0.0 <= beta <= 1.0
|
|
for p_ema, p in zip(net_ema.parameters(), net.parameters()):
|
|
p_ema.copy_(p.lerp(p_ema, beta))
|
|
if copy_buffer:
|
|
for b_ema, b in zip(net_ema.buffers(), net.buffers()):
|
|
b_ema.copy_(b)
|
|
|
|
|
|
def parallel(func, args_list, num_workers=32, timeout=None):
|
|
assert isinstance(args_list, list)
|
|
if not isinstance(args_list[0], tuple):
|
|
args_list = [(args, ) for args in args_list]
|
|
if num_workers == 0:
|
|
return [func(*args) for args in args_list]
|
|
with Pool(processes=num_workers) as pool:
|
|
results = [pool.apply_async(func, args) for args in args_list]
|
|
results = [res.get(timeout=timeout) for res in results]
|
|
return results
|
|
|
|
|
|
def exists(filename):
|
|
if filename.startswith('oss://'):
|
|
bucket, path = parse_oss_url(filename)
|
|
return bucket.object_exists(path)
|
|
else:
|
|
return osp.exists(filename)
|
|
|
|
|
|
def download(url, filename=None, replace=False, quiet=False):
|
|
if filename is None:
|
|
filename = osp.basename(url)
|
|
if not osp.exists(filename) or replace:
|
|
try:
|
|
if url.startswith('oss://'):
|
|
bucket, oss_key = parse_oss_url(url)
|
|
bucket.get_object_to_file(oss_key, filename)
|
|
else:
|
|
urllib.request.urlretrieve(url, filename)
|
|
if not quiet:
|
|
print(f'Downloaded {url} to {filename}', flush=True)
|
|
except Exception as e:
|
|
raise ValueError(f'Downloading {filename} failed with error {e}')
|
|
return osp.abspath(filename)
|
|
|
|
|
|
def unzip(filename, dst_dir=None):
|
|
if dst_dir is None:
|
|
dst_dir = osp.dirname(filename)
|
|
with zipfile.ZipFile(filename, 'r') as zip_ref:
|
|
zip_ref.extractall(dst_dir)
|
|
|
|
|
|
def load_state_dict(module, state_dict, drop_prefix=''):
|
|
# find incompatible key-vals
|
|
src, dst = state_dict, module.state_dict()
|
|
if drop_prefix:
|
|
src = type(src)([
|
|
(k[len(drop_prefix):] if k.startswith(drop_prefix) else k, v)
|
|
for k, v in src.items()
|
|
])
|
|
missing = [k for k in dst if k not in src]
|
|
unexpected = [k for k in src if k not in dst]
|
|
unmatched = [
|
|
k for k in src.keys() & dst.keys() if src[k].shape != dst[k].shape
|
|
]
|
|
|
|
# keep only compatible key-vals
|
|
incompatible = set(unexpected + unmatched)
|
|
src = type(src)([(k, v) for k, v in src.items() if k not in incompatible])
|
|
module.load_state_dict(src, strict=False)
|
|
|
|
# report incompatible key-vals
|
|
if len(missing) != 0:
|
|
print(' Missing: ' + ', '.join(missing), flush=True)
|
|
if len(unexpected) != 0:
|
|
print(' Unexpected: ' + ', '.join(unexpected), flush=True)
|
|
if len(unmatched) != 0:
|
|
print(' Shape unmatched: ' + ', '.join(unmatched), flush=True)
|
|
|
|
|
|
def inverse_indices(indices):
|
|
r"""Inverse map of indices.
|
|
E.g., if A[indices] == B, then B[inv_indices] == A.
|
|
"""
|
|
inv_indices = torch.empty_like(indices)
|
|
inv_indices[indices] = torch.arange(len(indices)).to(indices)
|
|
return inv_indices
|
|
|
|
|
|
def detect_duplicates(feats, thr=0.9):
|
|
assert feats.ndim == 2
|
|
|
|
# compute simmat
|
|
feats = F.normalize(feats, p=2, dim=1)
|
|
simmat = torch.mm(feats, feats.T)
|
|
simmat.triu_(1)
|
|
torch.cuda.synchronize()
|
|
|
|
# detect duplicates
|
|
mask = ~simmat.gt(thr).any(dim=0)
|
|
return torch.where(mask)[0]
|
|
|
|
|
|
class TFSClient(object):
|
|
|
|
def __init__(self,
|
|
host='restful-store.vip.tbsite.net:3800',
|
|
app_key='5354c9fae75f5'):
|
|
self.host = host
|
|
self.app_key = app_key
|
|
|
|
# candidate servers
|
|
self.servers = [
|
|
u for u in read(f'http://{host}/url.list').strip().split('\n')[1:]
|
|
if ':' in u
|
|
]
|
|
assert len(self.servers) >= 1
|
|
self.__server_id = -1
|
|
|
|
@property
|
|
def server(self):
|
|
self.__server_id = (self.__server_id + 1) % len(self.servers)
|
|
return self.servers[self.__server_id]
|
|
|
|
def read(self, tfs):
|
|
tfs = osp.basename(tfs)
|
|
meta = json.loads(
|
|
read(
|
|
f'http://{self.server}/v1/{self.app_key}/metadata/{tfs}?force=0'
|
|
))
|
|
img = Image.open(
|
|
BytesIO(
|
|
read(
|
|
f'http://{self.server}/v1/{self.app_key}/{tfs}?offset=0&size={meta["SIZE"]}',
|
|
'rb')))
|
|
return img
|
|
|
|
|
|
def read_tfs(tfs, retry=5):
|
|
exception = None
|
|
for _ in range(retry):
|
|
try:
|
|
global TFS_CLIENT
|
|
if TFS_CLIENT is None:
|
|
TFS_CLIENT = TFSClient()
|
|
return TFS_CLIENT.read(tfs)
|
|
except Exception as e:
|
|
exception = e
|
|
continue
|
|
else:
|
|
raise exception
|
|
|
|
|
|
def md5(filename):
|
|
with open(filename, 'rb') as f:
|
|
return hashlib.md5(f.read()).hexdigest()
|
|
|
|
|
|
def rope(x):
|
|
r"""Apply rotary position embedding on x of shape [B, *(spatial dimensions), C].
|
|
"""
|
|
# reshape
|
|
shape = x.shape
|
|
x = x.view(x.size(0), -1, x.size(-1))
|
|
l, c = x.shape[-2:]
|
|
assert c % 2 == 0
|
|
half = c // 2
|
|
|
|
# apply rotary position embedding on x
|
|
sinusoid = torch.outer(
|
|
torch.arange(l).to(x),
|
|
torch.pow(10000, -torch.arange(half).to(x).div(half)))
|
|
sin, cos = torch.sin(sinusoid), torch.cos(sinusoid)
|
|
x1, x2 = x.chunk(2, dim=-1)
|
|
x = torch.cat([x1 * cos - x2 * sin, x2 * cos + x1 * sin], dim=-1)
|
|
|
|
# reshape back
|
|
return x.view(shape)
|
|
|
|
|
|
def format_state(state, filename=None):
|
|
r"""For comparing/aligning state_dict.
|
|
"""
|
|
content = '\n'.join([f'{k}\t{tuple(v.shape)}' for k, v in state.items()])
|
|
if filename:
|
|
with open(filename, 'w') as f:
|
|
f.write(content)
|
|
|
|
|
|
def breakup_grid(img, grid_size):
|
|
r"""The inverse operator of ``torchvision.utils.make_grid``.
|
|
"""
|
|
# params
|
|
nrow = img.height // grid_size
|
|
ncol = img.width // grid_size
|
|
wrow = wcol = 2
|
|
|
|
# collect grids
|
|
grids = []
|
|
for i in range(nrow):
|
|
for j in range(ncol):
|
|
x1 = j * grid_size + (j + 1) * wcol
|
|
y1 = i * grid_size + (i + 1) * wrow
|
|
grids.append(img.crop((x1, y1, x1 + grid_size, y1 + grid_size)))
|
|
return grids
|
|
|
|
|
|
def huggingface_tokenizer(name='google/mt5-xxl', **kwargs):
|
|
from transformers import AutoTokenizer
|
|
return AutoTokenizer.from_pretrained(
|
|
DOWNLOAD_TO_CACHE(f'huggingface/tokenizers/{name}', name), **kwargs)
|
|
|
|
|
|
def huggingface_model(name='google/mt5-xxl', model_type='AutoModel', **kwargs):
|
|
import transformers
|
|
return getattr(transformers, model_type).from_pretrained(
|
|
DOWNLOAD_TO_CACHE(f'huggingface/models/{name}', name), **kwargs)
|