import random
from typing import Dict, Optional, Tuple
import torch
import torch.nn as nn
import torch.nn.functional as F
import numpy as np
from omegaconf import DictConfig
from tools.torch_utils import persistence
from tools.torch_utils.ops import bias_act, upfirdn2d, conv2d_resample
from tools.torch_utils import misc
#----------------------------------------------------------------------------
@misc.profiled_function
def normalize_2nd_moment(x, dim=1, eps=1e-8):
return x * (x.square().mean(dim=dim, keepdim=True) + eps).rsqrt()
#----------------------------------------------------------------------------
@persistence.persistent_class
class MappingNetwork(torch.nn.Module):
def __init__(self,
z_dim, # Input latent (Z) dimensionality, 0 = no latent.
c_dim, # Conditioning label (C) dimensionality, 0 = no label.
w_dim, # Intermediate latent (W) dimensionality.
num_ws, # Number of intermediate latents to output, None = do not broadcast.
num_layers = 8, # Number of mapping layers.
embed_features = None, # Label embedding dimensionality, None = same as w_dim.
layer_features = None, # Number of intermediate features in the mapping layers, None = same as w_dim.
activation = 'lrelu', # Activation function: 'relu', 'lrelu', etc.
lr_multiplier = 0.01, # Learning rate multiplier for the mapping layers.
w_avg_beta = 0.995, # Decay for tracking the moving average of W during training, None = do not track.
cfg = {}, # Additional config
):
super().__init__()
self.cfg = cfg
self.z_dim = z_dim
self.c_dim = c_dim
self.w_dim = w_dim
self.num_ws = num_ws
self.num_layers = num_layers
self.w_avg_beta = w_avg_beta
if embed_features is None:
embed_features = w_dim
if c_dim == 0:
embed_features = 0
if layer_features is None:
layer_features = w_dim
features_list = [z_dim + embed_features] + [layer_features] * (num_layers - 1) + [w_dim]
if c_dim > 0:
self.embed = FullyConnectedLayer(c_dim, embed_features)
for idx in range(num_layers):
in_features = features_list[idx]
out_features = features_list[idx + 1]
layer = FullyConnectedLayer(in_features, out_features, activation=activation, lr_multiplier=lr_multiplier)
setattr(self, f'fc{idx}', layer)
if num_ws is not None and w_avg_beta is not None:
self.register_buffer('w_avg', torch.zeros([w_dim]))
def forward(self, z, c, truncation_psi=1, truncation_cutoff=None, skip_w_avg_update=False):
# Embed, normalize, and concat inputs.
x = None
with torch.autograd.profiler.record_function('input'):
if self.z_dim > 0:
misc.assert_shape(z, [None, self.z_dim])
x = normalize_2nd_moment(z.to(torch.float32))
if self.c_dim > 0:
misc.assert_shape(c, [None, self.c_dim])
y = normalize_2nd_moment(self.embed(c.to(torch.float32)))
x = torch.cat([x, y], dim=1) if x is not None else y
# Main layers.
for idx in range(self.num_layers):
layer = getattr(self, f'fc{idx}')
x = layer(x)
# Update moving average of W.
if self.w_avg_beta is not None and self.training and not skip_w_avg_update:
with torch.autograd.profiler.record_function('update_w_avg'):
self.w_avg.copy_(x.detach().mean(dim=0).lerp(self.w_avg, self.w_avg_beta))
# Broadcast.
if self.num_ws is not None:
with torch.autograd.profiler.record_function('broadcast'):
x = x.unsqueeze(1).repeat([1, self.num_ws, 1])
# Apply truncation.
if truncation_psi != 1:
with torch.autograd.profiler.record_function('truncate'):
assert self.w_avg_beta is not None
if self.num_ws is None or truncation_cutoff is None:
x = self.w_avg.lerp(x, truncation_psi)
else:
x[:, :truncation_cutoff] = self.w_avg.lerp(x[:, :truncation_cutoff], truncation_psi)
return x
#----------------------------------------------------------------------------
@persistence.persistent_class
class FullyConnectedLayer(torch.nn.Module):
def __init__(self,
in_features, # Number of input features.
out_features, # Number of output features.
bias = True, # Apply additive bias before the activation function?
activation = 'linear', # Activation function: 'relu', 'lrelu', etc.
lr_multiplier = 1, # Learning rate multiplier.
bias_init = 0, # Initial value for the additive bias.
):
super().__init__()
self.activation = activation
self.weight = torch.nn.Parameter(torch.randn([out_features, in_features]) / lr_multiplier)
self.bias = torch.nn.Parameter(torch.full([out_features], float(bias_init))) if bias else None
self.weight_gain = lr_multiplier / np.sqrt(in_features)
self.bias_gain = lr_multiplier
def forward(self, x):
w = self.weight.to(x.dtype) * self.weight_gain
b = self.bias
if b is not None:
b = b.to(x.dtype)
if self.bias_gain != 1:
b = b * self.bias_gain
if self.activation == 'linear' and b is not None:
x = torch.addmm(b.unsqueeze(0), x, w.t())
else:
x = x.matmul(w.t())
x = bias_act.bias_act(x, b, act=self.activation)
return x
#----------------------------------------------------------------------------
@persistence.persistent_class
class Conv2dLayer(torch.nn.Module):
def __init__(self,
in_channels, # Number of input channels.
out_channels, # Number of output channels.
kernel_size, # Width and height of the convolution kernel.
bias = True, # Apply additive bias before the activation function?
activation = 'linear', # Activation function: 'relu', 'lrelu', etc.
up = 1, # Integer upsampling factor.
down = 1, # Integer downsampling factor.
resample_filter = [1,3,3,1], # Low-pass filter to apply when resampling activations.
conv_clamp = None, # Clamp the output to +-X, None = disable clamping.
channels_last = False, # Expect the input to have memory_format=channels_last?
trainable = True, # Update the weights of this layer during training?
instance_norm = False, # Should we apply instance normalization to y?
lr_multiplier = 1.0, # Learning rate multiplier.
):
super().__init__()
self.activation = activation
self.up = up
self.down = down
self.conv_clamp = conv_clamp
self.register_buffer('resample_filter', upfirdn2d.setup_filter(resample_filter))
self.padding = kernel_size // 2
self.weight_gain = 1 / np.sqrt(in_channels * (kernel_size ** 2))
self.act_gain = bias_act.activation_funcs[activation].def_gain
self.instance_norm = instance_norm
self.lr_multiplier = lr_multiplier
memory_format = torch.channels_last if channels_last else torch.contiguous_format
weight = torch.randn([out_channels, in_channels, kernel_size, kernel_size]).to(memory_format=memory_format)
bias = torch.zeros([out_channels]) if bias else None
if trainable:
self.weight = torch.nn.Parameter(weight)
self.bias = torch.nn.Parameter(bias) if bias is not None else None
else:
self.register_buffer('weight', weight)
if bias is not None:
self.register_buffer('bias', bias)
else:
self.bias = None
def forward(self, x, gain=1):
w = self.weight * (self.weight_gain * self.lr_multiplier)
b = (self.bias.to(x.dtype) * self.lr_multiplier) if self.bias is not None else None
flip_weight = (self.up == 1) # slightly faster
x = conv2d_resample.conv2d_resample(x=x, w=w.to(x.dtype), f=self.resample_filter, up=self.up, down=self.down, padding=self.padding, flip_weight=flip_weight)
act_gain = self.act_gain * gain
act_clamp = self.conv_clamp * gain if self.conv_clamp is not None else None
x = bias_act.bias_act(x, b, act=self.activation, gain=act_gain, clamp=act_clamp)
if self.instance_norm:
x = (x - x.mean(dim=(2,3), keepdim=True)) / (x.std(dim=(2,3), keepdim=True) + 1e-8) # [batch_size, c, h, w]
return x
#----------------------------------------------------------------------------
@persistence.persistent_class
class GenInput(nn.Module):
def __init__(self, cfg: DictConfig, channel_dim: int, motion_v_dim: int=None):
super().__init__()
self.cfg = cfg
if self.cfg.input.type == 'const':
self.input = torch.nn.Parameter(torch.randn([channel_dim, 4, 4]))
self.total_dim = channel_dim
elif self.cfg.input.type == 'temporal':
self.input = TemporalInput(self.cfg, channel_dim, motion_v_dim=motion_v_dim)
self.total_dim = self.input.get_dim()
else:
raise NotImplementedError(f'Unkown input type: {self.cfg.input.type}')
def forward(self, batch_size: int, motion_v: Optional[torch.Tensor]=None, dtype=None, memory_format=None) -> torch.Tensor:
if self.cfg.input.type == 'const':
x = self.input.to(dtype=dtype, memory_format=memory_format)
x = x.unsqueeze(0).repeat([batch_size, 1, 1, 1])
elif self.cfg.input.type == 'temporal':
x = self.input(motion_v=motion_v) # [batch_size, d, h, w]
else:
raise NotImplementedError(f'Unkown input type: {self.cfg.input.type}')
return x
#----------------------------------------------------------------------------
@persistence.persistent_class
class TemporalInput(nn.Module):
def __init__(self, cfg: DictConfig, channel_dim: int, motion_v_dim: int):
super().__init__()
self.cfg = cfg
self.motion_v_dim = motion_v_dim
self.const = nn.Parameter(torch.randn(1, channel_dim, 4, 4))
def get_dim(self):
return self.motion_v_dim + self.const.shape[1]
def forward(self, motion_v: torch.Tensor) -> torch.Tensor:
"""
motion_v: [batch_size, motion_v_dim]
"""
out = torch.cat([
self.const.repeat(len(motion_v), 1, 1, 1),
motion_v.unsqueeze(2).unsqueeze(3).repeat(1, 1, *self.const.shape[2:]),
], dim=1) # [batch_size, channel_dim + num_fourier_feats * 2]
return out
#----------------------------------------------------------------------------
class TemporalDifferenceEncoder(nn.Module):
def __init__(self, cfg: DictConfig):
super().__init__()
self.cfg = cfg
if self.cfg.sampling.num_frames_per_video > 1:
self.d = 256
self.const_embed = nn.Embedding(self.cfg.sampling.max_num_frames, self.d)
self.time_encoder = FixedTimeEncoder(
self.cfg.sampling.max_num_frames,
skip_small_t_freqs=self.cfg.get('skip_small_t_freqs', 0))
def get_dim(self) -> int:
if self.cfg.sampling.num_frames_per_video == 1:
return 1
else:
if self.cfg.sampling.type == 'uniform':
return self.d + self.time_encoder.get_dim()
else:
return (self.d + self.time_encoder.get_dim()) * (self.cfg.sampling.num_frames_per_video - 1)
def forward(self, t: torch.Tensor) -> torch.Tensor:
misc.assert_shape(t, [None, self.cfg.sampling.num_frames_per_video])
batch_size = t.shape[0]
if self.cfg.sampling.num_frames_per_video == 1:
out = torch.zeros(len(t), 1, device=t.device)
else:
if self.cfg.sampling.type == 'uniform':
num_diffs_to_use = 1
t_diffs = t[:, 1] - t[:, 0] # [batch_size]
else:
num_diffs_to_use = self.cfg.sampling.num_frames_per_video - 1
t_diffs = (t[:, 1:] - t[:, :-1]).view(-1) # [batch_size * (num_frames - 1)]
# Note: float => round => long is necessary when it's originally long
const_embs = self.const_embed(t_diffs.float().round().long()) # [batch_size * num_diffs_to_use, d]
fourier_embs = self.time_encoder(t_diffs.unsqueeze(1)) # [batch_size * num_diffs_to_use, num_fourier_feats]
out = torch.cat([const_embs, fourier_embs], dim=1) # [batch_size * num_diffs_to_use, d + num_fourier_feats]
out = out.view(batch_size, num_diffs_to_use, -1).view(batch_size, -1) # [batch_size, num_diffs_to_use * (d + num_fourier_feats)]
return out
#----------------------------------------------------------------------------
@persistence.persistent_class
class FixedTimeEncoder(nn.Module):
def __init__(self,
max_num_frames: int, # Maximum T size
skip_small_t_freqs: int=0, # How many high frequencies we should skip
):
super().__init__()
assert max_num_frames >= 1, f"Wrong max_num_frames: {max_num_frames}"
fourier_coefs = construct_log_spaced_freqs(max_num_frames, skip_small_t_freqs=skip_small_t_freqs)
self.register_buffer('fourier_coefs', fourier_coefs) # [1, num_fourier_feats]
def get_dim(self) -> int:
return self.fourier_coefs.shape[1] * 2
def forward(self, t: torch.Tensor) -> torch.Tensor:
assert t.ndim == 2, f"Wrong shape: {t.shape}"
t = t.view(-1).float() # [batch_size * num_frames]
fourier_raw_embs = self.fourier_coefs * t.unsqueeze(1) # [bf, num_fourier_feats]
fourier_embs = torch.cat([
fourier_raw_embs.sin(),
fourier_raw_embs.cos(),
], dim=1) # [bf, num_fourier_feats * 2]
return fourier_embs
#----------------------------------------------------------------------------
@persistence.persistent_class
class EqLRConv1d(nn.Module):
def __init__(self,
in_features: int,
out_features: int,
kernel_size: int,
padding: int=0,
stride: int=1,
activation: str='linear',
lr_multiplier: float=1.0,
bias=True,
bias_init=0.0,
):
super().__init__()
self.activation = activation
self.weight = torch.nn.Parameter(torch.randn([out_features, in_features, kernel_size]) / lr_multiplier)
self.bias = torch.nn.Parameter(torch.full([out_features], float(bias_init))) if bias else None
self.weight_gain = lr_multiplier / np.sqrt(in_features * kernel_size)
self.bias_gain = lr_multiplier
self.padding = padding
self.stride = stride
assert self.activation in ['lrelu', 'linear']
def forward(self, x: torch.Tensor) -> torch.Tensor:
assert x.ndim == 3, f"Wrong shape: {x.shape}"
w = self.weight.to(x.dtype) * self.weight_gain # [out_features, in_features, kernel_size]
b = self.bias # [out_features]
if b is not None:
b = b.to(x.dtype)
if self.bias_gain != 1:
b = b * self.bias_gain
y = F.conv1d(input=x, weight=w, bias=b, stride=self.stride, padding=self.padding) # [batch_size, out_features, out_len]
if self.activation == 'linear':
pass
elif self.activation == 'lrelu':
y = F.leaky_relu(y, negative_slope=0.2) # [batch_size, out_features, out_len]
else:
raise NotImplementedError
return y
#----------------------------------------------------------------------------
def sample_frames(cfg: Dict, total_video_len: int, **kwargs) -> np.ndarray:
if cfg['type'] == 'random':
return random_frame_sampling(cfg, total_video_len, **kwargs)
elif cfg['type'] == 'uniform':
return uniform_frame_sampling(cfg, total_video_len, **kwargs)
else:
raise NotImplementedError
#----------------------------------------------------------------------------
def random_frame_sampling(cfg: Dict, total_video_len: int, use_fractional_t: bool=False) -> np.ndarray:
min_time_diff = cfg["num_frames_per_video"] - 1
max_time_diff = min(total_video_len - 1, cfg.get('max_dist', float('inf')))
if type(cfg.get('total_dists')) in (list, tuple):
time_diff_range = [d for d in cfg['total_dists'] if min_time_diff <= d <= max_time_diff]
else:
time_diff_range = range(min_time_diff, max_time_diff)
time_diff: int = random.choice(time_diff_range)
if use_fractional_t:
offset = random.random() * (total_video_len - time_diff - 1)
else:
offset = random.randint(0, total_video_len - time_diff - 1)
frames_idx = [offset]
if cfg["num_frames_per_video"] > 1:
frames_idx.append(offset + time_diff)
if cfg["num_frames_per_video"] > 2:
frames_idx.extend([(offset + t) for t in random.sample(range(1, time_diff), k=cfg["num_frames_per_video"] - 2)])
frames_idx = sorted(frames_idx)
return np.array(frames_idx)
#----------------------------------------------------------------------------
def uniform_frame_sampling(cfg: Dict, total_video_len: int, use_fractional_t: bool=False) -> np.ndarray:
# Step 1: Select the distance between frames
if type(cfg.get('dists_between_frames')) in (list, tuple):
valid_dists = [d for d in cfg['dists_between_frames'] if d <= ['max_dist_between_frames']]
valid_dists = [d for d in valid_dists if (d * cfg['num_frames_per_video'] - d + 1) <= total_video_len]
d = random.choice(valid_dists)
else:
max_dist = min(cfg.get('max_dist', float('inf')), total_video_len // cfg['num_frames_per_video'])
d = random.randint(1, max_dist)
d_total = d * cfg['num_frames_per_video'] - d + 1
# Step 2: Sample.
if use_fractional_t:
offset = random.random() * (total_video_len - d_total)
else:
offset = random.randint(0, total_video_len - d_total)
frames_idx = offset + np.arange(cfg['num_frames_per_video']) * d
return frames_idx
#----------------------------------------------------------------------------
def construct_log_spaced_freqs(max_num_frames: int, skip_small_t_freqs: int=0) -> Tuple[int, torch.Tensor]:
time_resolution = 2 ** np.ceil(np.log2(max_num_frames))
num_fourier_feats = np.ceil(np.log2(time_resolution)).astype(int)
powers = torch.tensor([2]).repeat(num_fourier_feats).pow(torch.arange(num_fourier_feats)) # [num_fourier_feats]
powers = powers[:len(powers) - skip_small_t_freqs] # [num_fourier_feats]
fourier_coefs = powers.unsqueeze(0).float() * np.pi # [1, num_fourier_feats]
return fourier_coefs / time_resolution
#----------------------------------------------------------------------------