File size: 6,036 Bytes
5b3329b 986404c 5b3329b 986404c 5b3329b 986404c 5b3329b 986404c 5b3329b 986404c 5b3329b 986404c 5b3329b 986404c 5b3329b 986404c 5b3329b 986404c 5b3329b 986404c 5b3329b 986404c 5b3329b 986404c 5b3329b 986404c 5b3329b 986404c 5b3329b 986404c 5b3329b 986404c 5b3329b 986404c 5b3329b 986404c 5b3329b 986404c 5b3329b | 1 2 3 4 5 6 7 8 9 10 11 12 13 14 15 16 17 18 19 20 21 22 23 24 25 26 27 28 29 30 31 32 33 34 35 36 37 38 39 40 41 42 43 44 45 46 47 48 49 50 51 52 53 54 55 56 57 58 59 60 61 62 63 64 65 66 67 68 69 70 71 72 73 74 75 76 77 78 79 80 81 82 83 84 85 86 87 88 89 90 91 92 93 94 95 96 97 98 99 100 101 102 103 104 105 106 107 108 109 110 111 112 113 114 115 116 117 118 119 120 121 122 123 124 125 126 127 128 129 130 131 132 133 134 135 136 137 138 139 140 141 142 143 144 | """Compact but faithful conditional-regression and residual-EDM CorrDiff."""
import math
import torch
from torch import nn
from torch.nn import functional as F
class ConvBlock(nn.Module):
def __init__(self, in_channels, out_channels):
super().__init__()
groups = min(8, out_channels)
self.block = nn.Sequential(
nn.Conv2d(in_channels, out_channels, 3, padding=1),
nn.GroupNorm(groups, out_channels),
nn.SiLU(),
nn.Conv2d(out_channels, out_channels, 3, padding=1),
nn.GroupNorm(groups, out_channels),
nn.SiLU(),
)
self.skip = nn.Conv2d(in_channels, out_channels, 1)
def forward(self, x):
return self.block(x) + self.skip(x)
class RegressionUNet(nn.Module):
"""Deterministic conditional mean, computed cheaply before 448 px recovery."""
def __init__(self, in_channels=12, out_channels=4, base_channels=16, feature_size=56,
output_size=448):
super().__init__()
self.feature_size = feature_size
self.output_size = output_size
self.net = nn.Sequential(
ConvBlock(in_channels, base_channels),
ConvBlock(base_channels, base_channels * 2),
ConvBlock(base_channels * 2, base_channels),
nn.Conv2d(base_channels, out_channels, 1),
)
def forward(self, coarse):
x = F.interpolate(coarse, (self.feature_size, self.feature_size), mode="bilinear",
align_corners=False)
x = self.net(x)
return F.interpolate(x, (self.output_size, self.output_size), mode="bilinear",
align_corners=False)
class ResidualDenoiser(nn.Module):
"""Noise-conditional network used inside EDM preconditioning."""
def __init__(self, condition_channels=16, out_channels=4, base_channels=16,
feature_size=56):
super().__init__()
self.feature_size = feature_size
self.noise_mlp = nn.Sequential(
nn.Linear(1, base_channels), nn.SiLU(), nn.Linear(base_channels, base_channels)
)
self.input = ConvBlock(condition_channels, base_channels)
self.body = nn.Sequential(
ConvBlock(base_channels, base_channels * 2),
ConvBlock(base_channels * 2, base_channels),
nn.Conv2d(base_channels, out_channels, 1),
)
def forward(self, noisy, condition, c_noise):
size = noisy.shape[-2:]
x = torch.cat((noisy, condition), dim=1)
x = F.interpolate(x, (self.feature_size, self.feature_size), mode="bilinear",
align_corners=False)
x = self.input(x)
x = x + self.noise_mlp(c_noise[:, None].float())[:, :, None, None].to(x.dtype)
x = self.body(x)
return F.interpolate(x, size, mode="bilinear", align_corners=False)
class CorrDiff(nn.Module):
def __init__(self, in_channels=12, out_channels=4, base_channels=16, feature_size=56,
output_size=448, sigma_data=0.5):
super().__init__()
self.sigma_data = sigma_data
self.output_size = output_size
self.regression = RegressionUNet(in_channels, out_channels, base_channels,
feature_size, output_size)
self.diffusion = ResidualDenoiser(in_channels + 2 * out_channels, out_channels,
base_channels, feature_size)
def mean(self, coarse):
return self.regression(coarse)
def condition(self, coarse, mean):
coarse = F.interpolate(coarse, mean.shape[-2:], mode="bilinear", align_corners=False)
return torch.cat((coarse, mean), dim=1)
def denoise(self, noisy_residual, coarse, mean, sigma):
sigma = sigma.reshape(-1, 1, 1, 1).to(noisy_residual.dtype)
sigma_data = self.sigma_data
c_skip = sigma_data**2 / (sigma.square() + sigma_data**2)
c_out = sigma * sigma_data / (sigma.square() + sigma_data**2).sqrt()
c_in = (sigma.square() + sigma_data**2).rsqrt()
c_noise = sigma.flatten().log() / 4
network = self.diffusion(c_in * noisy_residual, self.condition(coarse, mean), c_noise)
return c_skip * noisy_residual + c_out * network
@staticmethod
def karras_schedule(steps, sigma_min, sigma_max, rho, device):
ramp = torch.linspace(0, 1, steps, device=device)
maximum = sigma_max ** (1 / rho)
minimum = sigma_min ** (1 / rho)
sigmas = (maximum + ramp * (minimum - maximum)) ** rho
return torch.cat((sigmas, sigmas.new_zeros(1)))
def sample(self, coarse, steps=4, sigma_min=0.002, sigma_max=5.0, rho=7.0,
solver="heun"):
mean = self.mean(coarse)
sigmas = self.karras_schedule(steps, sigma_min, sigma_max, rho, coarse.device)
x = torch.randn_like(mean) * sigmas[0]
for index, (current, following) in enumerate(zip(sigmas[:-1], sigmas[1:])):
sigma = current.expand(coarse.shape[0])
denoised = self.denoise(x, coarse, mean, sigma)
derivative = (x - denoised) / current
proposal = x + (following - current) * derivative
if solver == "heun" and index < len(sigmas) - 2:
next_sigma = following.expand(coarse.shape[0])
next_denoised = self.denoise(proposal, coarse, mean, next_sigma)
next_derivative = (proposal - next_denoised) / following
x = x + (following - current) * (derivative + next_derivative) / 2
else:
x = proposal
return mean + x
def forward(self, coarse, mode="sample", mean=None, noisy=None, sigma=None,
**sample_options):
if mode == "mean":
return self.mean(coarse)
if mode == "denoise":
return self.denoise(noisy, coarse, mean, sigma)
return self.sample(coarse, **sample_options)
__all__ = ["CorrDiff", "RegressionUNet", "ResidualDenoiser"]
|