City
b7a293cc27
Quick and lazy, just ported the example code. Can't do much more without having the model arch. https://github.com/openai/consistencydecoder/issues/1
117 lines
4.2 KiB
Python
117 lines
4.2 KiB
Python
import math
|
|
import torch
|
|
|
|
def _extract_into_tensor(arr, timesteps, broadcast_shape):
|
|
# from: https://github.com/openai/guided-diffusion/blob/22e0df8183507e13a7813f8d38d51b072ca1e67c/guided_diffusion/gaussian_diffusion.py#L895 """
|
|
res = arr[timesteps.to(torch.int).cpu()].float().to(timesteps.device)
|
|
dims_to_append = len(broadcast_shape) - len(res.shape)
|
|
return res[(...,) + (None,) * dims_to_append]
|
|
|
|
def betas_for_alpha_bar(num_diffusion_timesteps, alpha_bar, max_beta=0.999):
|
|
# from: https://github.com/openai/guided-diffusion/blob/22e0df8183507e13a7813f8d38d51b072ca1e67c/guided_diffusion/gaussian_diffusion.py#L45
|
|
betas = []
|
|
for i in range(num_diffusion_timesteps):
|
|
t1 = i / num_diffusion_timesteps
|
|
t2 = (i + 1) / num_diffusion_timesteps
|
|
betas.append(min(1 - alpha_bar(t2) / alpha_bar(t1), max_beta))
|
|
return torch.tensor(betas)
|
|
|
|
class ConsistencyDecoder(torch.nn.Module):
|
|
# From https://github.com/openai/consistencydecoder
|
|
def __init__(self, model_path, device):
|
|
super().__init__()
|
|
self.ckpt = torch.jit.load(model_path, map_location=device)
|
|
self.n_distilled_steps = 64
|
|
|
|
sigma_data = 0.5
|
|
betas = betas_for_alpha_bar(
|
|
1024, lambda t: math.cos((t + 0.008) / 1.008 * math.pi / 2) ** 2
|
|
)
|
|
alphas = 1.0 - betas
|
|
alphas_cumprod = torch.cumprod(alphas, dim=0)
|
|
self.sqrt_alphas_cumprod = torch.sqrt(alphas_cumprod)
|
|
self.sqrt_one_minus_alphas_cumprod = torch.sqrt(1.0 - alphas_cumprod)
|
|
sqrt_recip_alphas_cumprod = torch.sqrt(1.0 / alphas_cumprod)
|
|
sigmas = torch.sqrt(1.0 / alphas_cumprod - 1)
|
|
self.c_skip = (
|
|
sqrt_recip_alphas_cumprod
|
|
* sigma_data**2
|
|
/ (sigmas**2 + sigma_data**2)
|
|
)
|
|
self.c_out = sigmas * sigma_data / (sigmas**2 + sigma_data**2) ** 0.5
|
|
self.c_in = sqrt_recip_alphas_cumprod / (sigmas**2 + sigma_data**2) ** 0.5
|
|
|
|
@staticmethod
|
|
def round_timesteps(timesteps, total_timesteps, n_distilled_steps, truncate_start=True):
|
|
with torch.no_grad():
|
|
space = torch.div(total_timesteps, n_distilled_steps, rounding_mode="floor")
|
|
rounded_timesteps = (
|
|
torch.div(timesteps, space, rounding_mode="floor") + 1
|
|
) * space
|
|
if truncate_start:
|
|
rounded_timesteps[rounded_timesteps == total_timesteps] -= space
|
|
else:
|
|
rounded_timesteps[rounded_timesteps == total_timesteps] -= space
|
|
rounded_timesteps[rounded_timesteps == 0] += space
|
|
return rounded_timesteps
|
|
|
|
@staticmethod
|
|
def ldm_transform_latent(z, extra_scale_factor=1):
|
|
channel_means = [0.38862467, 0.02253063, 0.07381133, -0.0171294]
|
|
channel_stds = [0.9654121, 1.0440036, 0.76147926, 0.77022034]
|
|
|
|
if len(z.shape) != 4:
|
|
raise ValueError()
|
|
|
|
z = z * 0.18215
|
|
channels = [z[:, i] for i in range(z.shape[1])]
|
|
|
|
channels = [
|
|
extra_scale_factor * (c - channel_means[i]) / channel_stds[i]
|
|
for i, c in enumerate(channels)
|
|
]
|
|
return torch.stack(channels, dim=1)
|
|
|
|
@torch.no_grad()
|
|
def decode(self, features: torch.Tensor, schedule=[1.0, 0.5]):
|
|
features = self.ldm_transform_latent(features)
|
|
ts = self.round_timesteps(
|
|
torch.arange(0, 1024),
|
|
1024,
|
|
self.n_distilled_steps,
|
|
truncate_start=False,
|
|
)
|
|
shape = (
|
|
features.size(0),
|
|
3,
|
|
8 * features.size(2),
|
|
8 * features.size(3),
|
|
)
|
|
x_start = torch.zeros(shape, device=features.device, dtype=features.dtype)
|
|
schedule_timesteps = [int((1024 - 1) * s) for s in schedule]
|
|
for i in schedule_timesteps:
|
|
t = ts[i].item()
|
|
t_ = torch.tensor([t] * features.shape[0]).to(features.device)
|
|
noise = torch.randn_like(x_start)
|
|
x_start = (
|
|
_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
|
|
)
|
|
c_in = _extract_into_tensor(self.c_in, t_, x_start.shape)
|
|
model_output = self.ckpt(c_in * x_start, t_, features=features)
|
|
B, C = x_start.shape[:2]
|
|
model_output, _ = torch.split(model_output, C, dim=1)
|
|
pred_xstart = (
|
|
_extract_into_tensor(self.c_out, t_, x_start.shape) * model_output
|
|
+ _extract_into_tensor(self.c_skip, t_, x_start.shape) * x_start
|
|
).clamp(-1, 1)
|
|
x_start = pred_xstart
|
|
return x_start
|
|
|
|
def encode(self, *args, **kwargs):
|
|
raise NotImplementedError("ConsistencyDecoder can't be used for encoding!")
|