| 79 | |
| 80 | |
| 81 | class ConsistencyDecoder: |
| 82 | def __init__(self, device="cuda:0", download_root=os.path.expanduser("~/.cache/clip")): |
| 83 | self.n_distilled_steps = 64 |
| 84 | download_target = _download( |
| 85 | "https://openaipublic.azureedge.net/diff-vae/c9cebd3132dd9c42936d803e33424145a748843c8f716c0814838bdc8a2fe7cb/decoder.pt", |
| 86 | download_root, |
| 87 | ) |
| 88 | self.ckpt = torch.jit.load(download_target).to(device) |
| 89 | self.device = device |
| 90 | sigma_data = 0.5 |
| 91 | betas = betas_for_alpha_bar(1024, lambda t: math.cos((t + 0.008) / 1.008 * math.pi / 2) ** 2).to(device) |
| 92 | alphas = 1.0 - betas |
| 93 | alphas_cumprod = torch.cumprod(alphas, dim=0) |
| 94 | self.sqrt_alphas_cumprod = torch.sqrt(alphas_cumprod) |
| 95 | self.sqrt_one_minus_alphas_cumprod = torch.sqrt(1.0 - alphas_cumprod) |
| 96 | sqrt_recip_alphas_cumprod = torch.sqrt(1.0 / alphas_cumprod) |
| 97 | sigmas = torch.sqrt(1.0 / alphas_cumprod - 1) |
| 98 | self.c_skip = sqrt_recip_alphas_cumprod * sigma_data**2 / (sigmas**2 + sigma_data**2) |
| 99 | self.c_out = sigmas * sigma_data / (sigmas**2 + sigma_data**2) ** 0.5 |
| 100 | self.c_in = sqrt_recip_alphas_cumprod / (sigmas**2 + sigma_data**2) ** 0.5 |
| 101 | |
| 102 | @staticmethod |
| 103 | def round_timesteps(timesteps, total_timesteps, n_distilled_steps, truncate_start=True): |
| 104 | with torch.no_grad(): |
| 105 | space = torch.div(total_timesteps, n_distilled_steps, rounding_mode="floor") |
| 106 | rounded_timesteps = (torch.div(timesteps, space, rounding_mode="floor") + 1) * space |
| 107 | if truncate_start: |
| 108 | rounded_timesteps[rounded_timesteps == total_timesteps] -= space |
| 109 | else: |
| 110 | rounded_timesteps[rounded_timesteps == total_timesteps] -= space |
| 111 | rounded_timesteps[rounded_timesteps == 0] += space |
| 112 | return rounded_timesteps |
| 113 | |
| 114 | @staticmethod |
| 115 | def ldm_transform_latent(z, extra_scale_factor=1): |
| 116 | channel_means = [0.38862467, 0.02253063, 0.07381133, -0.0171294] |
| 117 | channel_stds = [0.9654121, 1.0440036, 0.76147926, 0.77022034] |
| 118 | |
| 119 | if len(z.shape) != 4: |
| 120 | raise ValueError() |
| 121 | |
| 122 | z = z * 0.18215 |
| 123 | channels = [z[:, i] for i in range(z.shape[1])] |
| 124 | |
| 125 | channels = [extra_scale_factor * (c - channel_means[i]) / channel_stds[i] for i, c in enumerate(channels)] |
| 126 | return torch.stack(channels, dim=1) |
| 127 | |
| 128 | @torch.no_grad() |
| 129 | def __call__( |
| 130 | self, |
| 131 | features: torch.Tensor, |
| 132 | schedule=[1.0, 0.5], |
| 133 | generator=None, |
| 134 | ): |
| 135 | features = self.ldm_transform_latent(features) |
| 136 | ts = self.round_timesteps( |
| 137 | torch.arange(0, 1024), |
| 138 | 1024, |
no outgoing calls
no test coverage detected