|
|
import torch |
|
|
import torch.nn.functional as F |
|
|
from torchvision import transforms |
|
|
|
|
|
def calc_mean_std(feat, eps=1e-5): |
|
|
|
|
|
size = feat.size() |
|
|
|
|
|
N, C = size[:2] |
|
|
feat_var = feat.view(N, C, -1).var(dim=2) + eps |
|
|
if len(size) == 3: |
|
|
feat_std = feat_var.sqrt().view(N, C, 1) |
|
|
feat_mean = feat.view(N, C, -1).mean(dim=2).view(N, C, 1) |
|
|
else: |
|
|
feat_std = feat_var.sqrt().view(N, C, 1, 1) |
|
|
feat_mean = feat.view(N, C, -1).mean(dim=2).view(N, C, 1, 1) |
|
|
return feat_mean, feat_std |
|
|
|
|
|
|
|
|
def get_img(img, resolution=512): |
|
|
norm_mean = [0.5, 0.5, 0.5] |
|
|
norm_std = [0.5, 0.5, 0.5] |
|
|
transform = transforms.Compose([ |
|
|
transforms.Resize((resolution, resolution)), |
|
|
transforms.ToTensor(), |
|
|
transforms.Normalize(norm_mean, norm_std) |
|
|
]) |
|
|
img = transform(img) |
|
|
return img.unsqueeze(0) |
|
|
|
|
|
@torch.no_grad() |
|
|
def slerp(p0, p1, fract_mixing: float, adain=True): |
|
|
r""" Copied from lunarring/latentblending |
|
|
Helper function to correctly mix two random variables using spherical interpolation. |
|
|
The function will always cast up to float64 for sake of extra 4. |
|
|
Args: |
|
|
p0: |
|
|
First tensor for interpolation |
|
|
p1: |
|
|
Second tensor for interpolation |
|
|
fract_mixing: float |
|
|
Mixing coefficient of interval [0, 1]. |
|
|
0 will return in p0 |
|
|
1 will return in p1 |
|
|
0.x will return a mix between both preserving angular velocity. |
|
|
""" |
|
|
if p0.dtype == torch.float16: |
|
|
recast_to = 'fp16' |
|
|
else: |
|
|
recast_to = 'fp32' |
|
|
|
|
|
p0 = p0.double() |
|
|
p1 = p1.double() |
|
|
|
|
|
if adain: |
|
|
mean1, std1 = calc_mean_std(p0) |
|
|
mean2, std2 = calc_mean_std(p1) |
|
|
mean = mean1 * (1 - fract_mixing) + mean2 * fract_mixing |
|
|
std = std1 * (1 - fract_mixing) + std2 * fract_mixing |
|
|
|
|
|
norm = torch.linalg.norm(p0) * torch.linalg.norm(p1) |
|
|
epsilon = 1e-7 |
|
|
dot = torch.sum(p0 * p1) / norm |
|
|
dot = dot.clamp(-1+epsilon, 1-epsilon) |
|
|
|
|
|
theta_0 = torch.arccos(dot) |
|
|
sin_theta_0 = torch.sin(theta_0) |
|
|
theta_t = theta_0 * fract_mixing |
|
|
s0 = torch.sin(theta_0 - theta_t) / sin_theta_0 |
|
|
s1 = torch.sin(theta_t) / sin_theta_0 |
|
|
interp = p0*s0 + p1*s1 |
|
|
|
|
|
if adain: |
|
|
interp = F.instance_norm(interp) * std + mean |
|
|
|
|
|
if recast_to == 'fp16': |
|
|
interp = interp.half() |
|
|
elif recast_to == 'fp32': |
|
|
interp = interp.float() |
|
|
|
|
|
return interp |
|
|
|
|
|
|
|
|
def do_replace_attn(key: str): |
|
|
|
|
|
return key.startswith('up') |
|
|
|
