commit 46d6daac580651313a2ed0a14376adfbd0b640fc Author: shadowcz007 Date: Fri Feb 2 09:35:36 2024 +0800 init diff --git a/__init__.py b/__init__.py new file mode 100644 index 0000000..feab1a7 --- /dev/null +++ b/__init__.py @@ -0,0 +1,10 @@ +# 导入节点 +from .nodes.Vae import VAELoader,VAEDecode + + +# 要导出的所有节点及其名称的字典 +# 注意：名称应全局唯一 +NODE_CLASS_MAPPINGS = { + "VAELoaderConsistencyDecoder":VAELoader, + "VAEDecodeConsistencyDecoder":VAEDecode, +} diff --git a/nodes/Vae.py b/nodes/Vae.py new file mode 100644 index 0000000..7b23634 --- /dev/null +++ b/nodes/Vae.py @@ -0,0 +1,178 @@ +# https://github.com/openai/consistencydecoder/blob/main/consistencydecoder/__init__.py + +import folder_paths +# from comfy import model_management +import math +import torch +import numpy as np +from PIL import Image + + +class ConsistencyDecoderWrapper: + def __init__(self, decoder): + self.decoder = decoder + def decode(self, x): + return self.decoder(x) + +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].float() + 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: + def __init__(self, device="cuda:0", download_target=""): + self.n_distilled_steps = 64 + # download_target = _download("https://openaipublic.azureedge.net/diff-vae/c9cebd3132dd9c42936d803e33424145a748843c8f716c0814838bdc8a2fe7cb/decoder.pt", download_root) + self.ckpt = torch.jit.load(download_target).to(device) + self.device = device + sigma_data = 0.5 + betas = betas_for_alpha_bar( + 1024, lambda t: math.cos((t + 0.008) / 1.008 * math.pi / 2) ** 2 + ).to(device) + 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 __call__( + 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(self.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 + + + +class VAELoader: + @classmethod + def INPUT_TYPES(s): + return {"required": { "vae_name": (folder_paths.get_filename_list("vae"), )}} + RETURN_TYPES = ("VAE",) + FUNCTION = "load_vae" + + CATEGORY = "openai/consistencydecoder" + + #TODO: scale factor? + def load_vae(self, vae_name): + vae_path = folder_paths.get_full_path("vae", vae_name) + device = 'cuda:0' + # print('device',device) + consistencyDecoder = ConsistencyDecoder(device=device, + download_target=vae_path) # Model size: 2.49 GB + vae = ConsistencyDecoderWrapper(consistencyDecoder) + return (vae,) + + +class VAEDecode: + @classmethod + def INPUT_TYPES(s): + return {"required": { "samples": ("LATENT", ), "vae": ("VAE", )}} + RETURN_TYPES = ("IMAGE",) + FUNCTION = "decode" + + CATEGORY = "openai/consistencydecoder" + + def decode(self, vae, samples): + image = vae.decode(samples["samples"].to("cuda:0")) + image = image[0].cpu().numpy() + image = (image + 1.0) * 127.5 + image = image.clip(0, 255).astype(np.uint8) + image = Image.fromarray(image.transpose(1, 2, 0)) + image = image.convert("RGB") + image = np.array(image).astype(np.float32) / 255.0 + image = torch.from_numpy(image)[None,] + return (image, ) \ No newline at end of file