194 lines
6.2 KiB
Python
194 lines
6.2 KiB
Python
from dataclasses import dataclass
|
|
from functools import partial
|
|
import cv2
|
|
import pandas as pd
|
|
import gc
|
|
import math
|
|
import lpips
|
|
from PIL import Image, ImageOps
|
|
import requests
|
|
import torch
|
|
from torch import nn
|
|
from torch.nn import functional as F
|
|
import torchvision
|
|
import torchvision.transforms as T
|
|
import torchvision.transforms.functional as TF
|
|
from tqdm import tqdm
|
|
from resize_right import resize
|
|
from guided_diffusion.script_util import create_model_and_diffusion, model_and_diffusion_defaults
|
|
import numpy as np
|
|
from functools import partial
|
|
from numpy import asarray
|
|
|
|
|
|
def append_dims(x, n):
|
|
return x[(Ellipsis, *(None,) * (n - x.ndim))]
|
|
|
|
|
|
def alpha_sigma_to_t(alpha, sigma):
|
|
return torch.atan2(sigma, alpha) * 2 / math.pi
|
|
|
|
|
|
def expand_to_planes(x, shape):
|
|
return append_dims(x, len(shape)).repeat([1, 1, *shape[2:]])
|
|
|
|
|
|
def t_to_alpha_sigma(t):
|
|
return torch.cos(t * math.pi / 2), torch.sin(t * math.pi / 2)
|
|
|
|
|
|
@dataclass
|
|
class DiffusionOutput:
|
|
v: torch.Tensor
|
|
pred: torch.Tensor
|
|
eps: torch.Tensor
|
|
|
|
|
|
class ConvBlock(nn.Sequential):
|
|
def __init__(self, c_in, c_out):
|
|
super().__init__(
|
|
nn.Conv2d(c_in, c_out, 3, padding=1),
|
|
nn.ReLU(inplace=True),
|
|
)
|
|
|
|
|
|
class SkipBlock(nn.Module):
|
|
def __init__(self, main, skip=None):
|
|
super().__init__()
|
|
self.main = nn.Sequential(*main)
|
|
self.skip = skip if skip else nn.Identity()
|
|
|
|
def forward(self, input):
|
|
return torch.cat([self.main(input), self.skip(input)], dim=1)
|
|
|
|
|
|
class FourierFeatures(nn.Module):
|
|
def __init__(self, in_features, out_features, std=1.):
|
|
super().__init__()
|
|
assert out_features % 2 == 0
|
|
self.weight = nn.Parameter(torch.randn(
|
|
[out_features // 2, in_features]) * std)
|
|
|
|
def forward(self, input):
|
|
f = 2 * math.pi * input @ self.weight.T
|
|
return torch.cat([f.cos(), f.sin()], dim=-1)
|
|
|
|
|
|
class SecondaryDiffusionImageNet(nn.Module):
|
|
def __init__(self):
|
|
super().__init__()
|
|
c = 64 # The base channel count
|
|
|
|
self.timestep_embed = FourierFeatures(1, 16)
|
|
|
|
self.net = nn.Sequential(
|
|
ConvBlock(3 + 16, c),
|
|
ConvBlock(c, c),
|
|
SkipBlock([
|
|
nn.AvgPool2d(2),
|
|
ConvBlock(c, c * 2),
|
|
ConvBlock(c * 2, c * 2),
|
|
SkipBlock([
|
|
nn.AvgPool2d(2),
|
|
ConvBlock(c * 2, c * 4),
|
|
ConvBlock(c * 4, c * 4),
|
|
SkipBlock([
|
|
nn.AvgPool2d(2),
|
|
ConvBlock(c * 4, c * 8),
|
|
ConvBlock(c * 8, c * 4),
|
|
nn.Upsample(scale_factor=2, mode='bilinear',
|
|
align_corners=False),
|
|
]),
|
|
ConvBlock(c * 8, c * 4),
|
|
ConvBlock(c * 4, c * 2),
|
|
nn.Upsample(scale_factor=2, mode='bilinear',
|
|
align_corners=False),
|
|
]),
|
|
ConvBlock(c * 4, c * 2),
|
|
ConvBlock(c * 2, c),
|
|
nn.Upsample(scale_factor=2, mode='bilinear',
|
|
align_corners=False),
|
|
]),
|
|
ConvBlock(c * 2, c),
|
|
nn.Conv2d(c, 3, 3, padding=1),
|
|
)
|
|
|
|
def forward(self, input, t):
|
|
timestep_embed = expand_to_planes(
|
|
self.timestep_embed(t[:, None]), input.shape)
|
|
v = self.net(torch.cat([input, timestep_embed], dim=1))
|
|
alphas, sigmas = map(
|
|
partial(append_dims, n=v.ndim), t_to_alpha_sigma(t))
|
|
pred = input * alphas - v * sigmas
|
|
eps = input * sigmas + v * alphas
|
|
return DiffusionOutput(v, pred, eps)
|
|
|
|
|
|
class SecondaryDiffusionImageNet2(nn.Module):
|
|
def __init__(self):
|
|
super().__init__()
|
|
c = 64 # The base channel count
|
|
cs = [c, c * 2, c * 2, c * 4, c * 4, c * 8]
|
|
|
|
self.timestep_embed = FourierFeatures(1, 16)
|
|
self.down = nn.AvgPool2d(2)
|
|
self.up = nn.Upsample(
|
|
scale_factor=2, mode='bilinear', align_corners=False)
|
|
|
|
self.net = nn.Sequential(
|
|
ConvBlock(3 + 16, cs[0]),
|
|
ConvBlock(cs[0], cs[0]),
|
|
SkipBlock([
|
|
self.down,
|
|
ConvBlock(cs[0], cs[1]),
|
|
ConvBlock(cs[1], cs[1]),
|
|
SkipBlock([
|
|
self.down,
|
|
ConvBlock(cs[1], cs[2]),
|
|
ConvBlock(cs[2], cs[2]),
|
|
SkipBlock([
|
|
self.down,
|
|
ConvBlock(cs[2], cs[3]),
|
|
ConvBlock(cs[3], cs[3]),
|
|
SkipBlock([
|
|
self.down,
|
|
ConvBlock(cs[3], cs[4]),
|
|
ConvBlock(cs[4], cs[4]),
|
|
SkipBlock([
|
|
self.down,
|
|
ConvBlock(cs[4], cs[5]),
|
|
ConvBlock(cs[5], cs[5]),
|
|
ConvBlock(cs[5], cs[5]),
|
|
ConvBlock(cs[5], cs[4]),
|
|
self.up,
|
|
]),
|
|
ConvBlock(cs[4] * 2, cs[4]),
|
|
ConvBlock(cs[4], cs[3]),
|
|
self.up,
|
|
]),
|
|
ConvBlock(cs[3] * 2, cs[3]),
|
|
ConvBlock(cs[3], cs[2]),
|
|
self.up,
|
|
]),
|
|
ConvBlock(cs[2] * 2, cs[2]),
|
|
ConvBlock(cs[2], cs[1]),
|
|
self.up,
|
|
]),
|
|
ConvBlock(cs[1] * 2, cs[1]),
|
|
ConvBlock(cs[1], cs[0]),
|
|
self.up,
|
|
]),
|
|
ConvBlock(cs[0] * 2, cs[0]),
|
|
nn.Conv2d(cs[0], 3, 3, padding=1),
|
|
)
|
|
|
|
def forward(self, input, t):
|
|
timestep_embed = expand_to_planes(
|
|
self.timestep_embed(t[:, None]), input.shape)
|
|
v = self.net(torch.cat([input, timestep_embed], dim=1))
|
|
alphas, sigmas = map(
|
|
partial(append_dims, n=v.ndim), t_to_alpha_sigma(t))
|
|
pred = input * alphas - v * sigmas
|
|
eps = input * sigmas + v * alphas
|
|
return DiffusionOutput(v, pred, eps)
|