Spaces:
Running
on
Zero
Running
on
Zero
File size: 4,230 Bytes
1de8821 |
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 |
from functools import partial
from typing import Tuple
import torch
from torch import nn
import numpy as np
def make_beta_schedule(schedule, n_timestep, linear_start=1e-4, linear_end=2e-2, cosine_s=8e-3):
if schedule == "linear":
betas = (
np.linspace(linear_start ** 0.5, linear_end ** 0.5, n_timestep, dtype=np.float64) ** 2
)
elif schedule == "cosine":
timesteps = (
np.arange(n_timestep + 1, dtype=np.float64) / n_timestep + cosine_s
)
alphas = timesteps / (1 + cosine_s) * np.pi / 2
alphas = np.cos(alphas).pow(2)
alphas = alphas / alphas[0]
betas = 1 - alphas[1:] / alphas[:-1]
betas = np.clip(betas, a_min=0, a_max=0.999)
elif schedule == "sqrt_linear":
betas = np.linspace(linear_start, linear_end, n_timestep, dtype=np.float64)
elif schedule == "sqrt":
betas = np.linspace(linear_start, linear_end, n_timestep, dtype=np.float64) ** 0.5
else:
raise ValueError(f"schedule '{schedule}' unknown.")
return betas
def extract_into_tensor(a: torch.Tensor, t: torch.Tensor, x_shape: Tuple[int]) -> torch.Tensor:
b, *_ = t.shape
out = a.gather(-1, t)
return out.reshape(b, *((1,) * (len(x_shape) - 1)))
class Diffusion(nn.Module):
def __init__(
self,
timesteps=1000,
beta_schedule="linear",
loss_type="l2",
linear_start=1e-4,
linear_end=2e-2,
cosine_s=8e-3,
parameterization="eps"
):
super().__init__()
self.num_timesteps = timesteps
self.beta_schedule = beta_schedule
self.linear_start = linear_start
self.linear_end = linear_end
self.cosine_s = cosine_s
assert parameterization in ["eps", "x0", "v"], "currently only supporting 'eps' and 'x0' and 'v'"
self.parameterization = parameterization
self.loss_type = loss_type
betas = make_beta_schedule(beta_schedule, timesteps, linear_start=linear_start, linear_end=linear_end,
cosine_s=cosine_s)
alphas = 1. - betas
alphas_cumprod = np.cumprod(alphas, axis=0)
sqrt_alphas_cumprod = np.sqrt(alphas_cumprod)
sqrt_one_minus_alphas_cumprod = np.sqrt(1. - alphas_cumprod)
self.betas = betas
self.register("sqrt_alphas_cumprod", sqrt_alphas_cumprod)
self.register("sqrt_one_minus_alphas_cumprod", sqrt_one_minus_alphas_cumprod)
def register(self, name: str, value: np.ndarray) -> None:
self.register_buffer(name, torch.tensor(value, dtype=torch.float32))
def q_sample(self, x_start, t, noise):
return (
extract_into_tensor(self.sqrt_alphas_cumprod, t, x_start.shape) * x_start +
extract_into_tensor(self.sqrt_one_minus_alphas_cumprod, t, x_start.shape) * noise
)
def get_v(self, x, noise, t):
return (
extract_into_tensor(self.sqrt_alphas_cumprod, t, x.shape) * noise -
extract_into_tensor(self.sqrt_one_minus_alphas_cumprod, t, x.shape) * x
)
def get_loss(self, pred, target, mean=True):
if self.loss_type == 'l1':
loss = (target - pred).abs()
if mean:
loss = loss.mean()
elif self.loss_type == 'l2':
if mean:
loss = torch.nn.functional.mse_loss(target, pred)
else:
loss = torch.nn.functional.mse_loss(target, pred, reduction='none')
else:
raise NotImplementedError("unknown loss type '{loss_type}'")
return loss
def p_losses(self, model, x_start, t, cond):
noise = torch.randn_like(x_start)
x_noisy = self.q_sample(x_start=x_start, t=t, noise=noise)
model_output = model(x_noisy, t, cond)
if self.parameterization == "x0":
target = x_start
elif self.parameterization == "eps":
target = noise
elif self.parameterization == "v":
target = self.get_v(x_start, noise, t)
else:
raise NotImplementedError()
loss_simple = self.get_loss(model_output, target, mean=False).mean()
return loss_simple
|