import math import os import urllib import warnings from argparse import ArgumentParser import torch import torch.nn as nn import torch.nn.functional as F from huggingface_hub.utils import insecure_hashlib from safetensors.torch import load_file as stl from tqdm import tqdm from diffusers import AutoencoderKL, ConsistencyDecoderVAE, DiffusionPipeline, StableDiffusionPipeline, UNet2DModel from diffusers.models.autoencoders.vae import Encoder from diffusers.models.embeddings import TimestepEmbedding from diffusers.models.unets.unet_2d_blocks import ResnetDownsampleBlock2D, ResnetUpsampleBlock2D, UNetMidBlock2D args = ArgumentParser() args.add_argument("--save_pretrained", required=False, default=None, type=str) args.add_argument("--test_image", required=True, type=str) args = args.parse_args() 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) def _download(url: str, root: str): os.makedirs(root, exist_ok=True) filename = os.path.basename(url) expected_sha256 = url.split("/")[-2] download_target = os.path.join(root, filename) if os.path.exists(download_target) and not os.path.isfile(download_target): raise RuntimeError(f"{download_target} exists and is not a regular file") if os.path.isfile(download_target): if insecure_hashlib.sha256(open(download_target, "rb").read()).hexdigest() == expected_sha256: return download_target else: warnings.warn(f"{download_target} exists, but the SHA256 checksum does not match; re-downloading the file") with urllib.request.urlopen(url) as source, open(download_target, "wb") as output: with tqdm( total=int(source.info().get("Content-Length")), ncols=80, unit="iB", unit_scale=True, unit_divisor=1024, ) as loop: while True: buffer = source.read(8192) if not buffer: break output.write(buffer) loop.update(len(buffer)) if insecure_hashlib.sha256(open(download_target, "rb").read()).hexdigest() != expected_sha256: raise RuntimeError("Model has been downloaded but the SHA256 checksum does not not match") return download_target class ConsistencyDecoder: def __init__(self, device="cuda:0", download_root=os.path.expanduser("~/.cache/clip")): 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], generator=None, ): 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) noise = torch.randn(x_start.shape, dtype=x_start.dtype, generator=generator).to(device=x_start.device) 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) import torch.nn.functional as F from diffusers import UNet2DModel if isinstance(self.ckpt, UNet2DModel): input = torch.concat([c_in * x_start, F.upsample_nearest(features, scale_factor=8)], dim=1) model_output = self.ckpt(input, t_).sample else: 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 save_image(image, name): import numpy as np from PIL import Image 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.save(name) def load_image(uri, size=None, center_crop=False): import numpy as np from PIL import Image image = Image.open(uri) if center_crop: image = image.crop( ( (image.width - min(image.width, image.height)) // 2, (image.height - min(image.width, image.height)) // 2, (image.width + min(image.width, image.height)) // 2, (image.height + min(image.width, image.height)) // 2, ) ) if size is not None: image = image.resize(size) image = torch.tensor(np.array(image).transpose(2, 0, 1)).unsqueeze(0).float() image = image / 127.5 - 1.0 return image class TimestepEmbedding_(nn.Module): def __init__(self, n_time=1024, n_emb=320, n_out=1280) -> None: super().__init__() self.emb = nn.Embedding(n_time, n_emb) self.f_1 = nn.Linear(n_emb, n_out) self.f_2 = nn.Linear(n_out, n_out) def forward(self, x) -> torch.Tensor: x = self.emb(x) x = self.f_1(x) x = F.silu(x) return self.f_2(x) class ImageEmbedding(nn.Module): def __init__(self, in_channels=7, out_channels=320) -> None: super().__init__() self.f = nn.Conv2d(in_channels, out_channels, kernel_size=3, padding=1) def forward(self, x) -> torch.Tensor: return self.f(x) class ImageUnembedding(nn.Module): def __init__(self, in_channels=320, out_channels=6) -> None: super().__init__() self.gn = nn.GroupNorm(32, in_channels) self.f = nn.Conv2d(in_channels, out_channels, kernel_size=3, padding=1) def forward(self, x) -> torch.Tensor: return self.f(F.silu(self.gn(x))) class ConvResblock(nn.Module): def __init__(self, in_features=320, out_features=320) -> None: super().__init__() self.f_t = nn.Linear(1280, out_features * 2) self.gn_1 = nn.GroupNorm(32, in_features) self.f_1 = nn.Conv2d(in_features, out_features, kernel_size=3, padding=1) self.gn_2 = nn.GroupNorm(32, out_features) self.f_2 = nn.Conv2d(out_features, out_features, kernel_size=3, padding=1) skip_conv = in_features != out_features self.f_s = nn.Conv2d(in_features, out_features, kernel_size=1, padding=0) if skip_conv else nn.Identity() def forward(self, x, t): x_skip = x t = self.f_t(F.silu(t)) t = t.chunk(2, dim=1) t_1 = t[0].unsqueeze(dim=2).unsqueeze(dim=3) + 1 t_2 = t[1].unsqueeze(dim=2).unsqueeze(dim=3) gn_1 = F.silu(self.gn_1(x)) f_1 = self.f_1(gn_1) gn_2 = self.gn_2(f_1) return self.f_s(x_skip) + self.f_2(F.silu(gn_2 * t_1 + t_2)) # Also ConvResblock class Downsample(nn.Module): def __init__(self, in_channels=320) -> None: super().__init__() self.f_t = nn.Linear(1280, in_channels * 2) self.gn_1 = nn.GroupNorm(32, in_channels) self.f_1 = nn.Conv2d(in_channels, in_channels, kernel_size=3, padding=1) self.gn_2 = nn.GroupNorm(32, in_channels) self.f_2 = nn.Conv2d(in_channels, in_channels, kernel_size=3, padding=1) def forward(self, x, t) -> torch.Tensor: x_skip = x t = self.f_t(F.silu(t)) t_1, t_2 = t.chunk(2, dim=1) t_1 = t_1.unsqueeze(2).unsqueeze(3) + 1 t_2 = t_2.unsqueeze(2).unsqueeze(3) gn_1 = F.silu(self.gn_1(x)) avg_pool2d = F.avg_pool2d(gn_1, kernel_size=(2, 2), stride=None) f_1 = self.f_1(avg_pool2d) gn_2 = self.gn_2(f_1) f_2 = self.f_2(F.silu(t_2 + (t_1 * gn_2))) return f_2 + F.avg_pool2d(x_skip, kernel_size=(2, 2), stride=None) # Also ConvResblock class Upsample(nn.Module): def __init__(self, in_channels=1024) -> None: super().__init__() self.f_t = nn.Linear(1280, in_channels * 2) self.gn_1 = nn.GroupNorm(32, in_channels) self.f_1 = nn.Conv2d(in_channels, in_channels, kernel_size=3, padding=1) self.gn_2 = nn.GroupNorm(32, in_channels) self.f_2 = nn.Conv2d(in_channels, in_channels, kernel_size=3, padding=1) def forward(self, x, t) -> torch.Tensor: x_skip = x t = self.f_t(F.silu(t)) t_1, t_2 = t.chunk(2, dim=1) t_1 = t_1.unsqueeze(2).unsqueeze(3) + 1 t_2 = t_2.unsqueeze(2).unsqueeze(3) gn_1 = F.silu(self.gn_1(x)) upsample = F.upsample_nearest(gn_1, scale_factor=2) f_1 = self.f_1(upsample) gn_2 = self.gn_2(f_1) f_2 = self.f_2(F.silu(t_2 + (t_1 * gn_2))) return f_2 + F.upsample_nearest(x_skip, scale_factor=2) class ConvUNetVAE(nn.Module): def __init__(self) -> None: super().__init__() self.embed_image = ImageEmbedding() self.embed_time = TimestepEmbedding_() down_0 = nn.ModuleList( [ ConvResblock(320, 320), ConvResblock(320, 320), ConvResblock(320, 320), Downsample(320), ] ) down_1 = nn.ModuleList( [ ConvResblock(320, 640), ConvResblock(640, 640), ConvResblock(640, 640), Downsample(640), ] ) down_2 = nn.ModuleList( [ ConvResblock(640, 1024), ConvResblock(1024, 1024), ConvResblock(1024, 1024), Downsample(1024), ] ) down_3 = nn.ModuleList( [ ConvResblock(1024, 1024), ConvResblock(1024, 1024), ConvResblock(1024, 1024), ] ) self.down = nn.ModuleList( [ down_0, down_1, down_2, down_3, ] ) self.mid = nn.ModuleList( [ ConvResblock(1024, 1024), ConvResblock(1024, 1024), ] ) up_3 = nn.ModuleList( [ ConvResblock(1024 * 2, 1024), ConvResblock(1024 * 2, 1024), ConvResblock(1024 * 2, 1024), ConvResblock(1024 * 2, 1024), Upsample(1024), ] ) up_2 = nn.ModuleList( [ ConvResblock(1024 * 2, 1024), ConvResblock(1024 * 2, 1024), ConvResblock(1024 * 2, 1024), ConvResblock(1024 + 640, 1024), Upsample(1024), ] ) up_1 = nn.ModuleList( [ ConvResblock(1024 + 640, 640), ConvResblock(640 * 2, 640), ConvResblock(640 * 2, 640), ConvResblock(320 + 640, 640), Upsample(640), ] ) up_0 = nn.ModuleList( [ ConvResblock(320 + 640, 320), ConvResblock(320 * 2, 320), ConvResblock(320 * 2, 320), ConvResblock(320 * 2, 320), ] ) self.up = nn.ModuleList( [ up_0, up_1, up_2, up_3, ] ) self.output = ImageUnembedding() def forward(self, x, t, features) -> torch.Tensor: converted = hasattr(self, "converted") and self.converted x = torch.cat([x, F.upsample_nearest(features, scale_factor=8)], dim=1) if converted: t = self.time_embedding(self.time_proj(t)) else: t = self.embed_time(t) x = self.embed_image(x) skips = [x] for i, down in enumerate(self.down): if converted and i in [0, 1, 2, 3]: x, skips_ = down(x, t) for skip in skips_: skips.append(skip) else: for block in down: x = block(x, t) skips.append(x) print(x.float().abs().sum()) if converted: x = self.mid(x, t) else: for i in range(2): x = self.mid[i](x, t) print(x.float().abs().sum()) for i, up in enumerate(self.up[::-1]): if converted and i in [0, 1, 2, 3]: skip_4 = skips.pop() skip_3 = skips.pop() skip_2 = skips.pop() skip_1 = skips.pop() skips_ = (skip_1, skip_2, skip_3, skip_4) x = up(x, skips_, t) else: for block in up: if isinstance(block, ConvResblock): x = torch.concat([x, skips.pop()], dim=1) x = block(x, t) return self.output(x) def rename_state_dict_key(k): k = k.replace("blocks.", "") for i in range(5): k = k.replace(f"down_{i}_", f"down.{i}.") k = k.replace(f"conv_{i}.", f"{i}.") k = k.replace(f"up_{i}_", f"up.{i}.") k = k.replace(f"mid_{i}", f"mid.{i}") k = k.replace("upsamp.", "4.") k = k.replace("downsamp.", "3.") k = k.replace("f_t.w", "f_t.weight").replace("f_t.b", "f_t.bias") k = k.replace("f_1.w", "f_1.weight").replace("f_1.b", "f_1.bias") k = k.replace("f_2.w", "f_2.weight").replace("f_2.b", "f_2.bias") k = k.replace("f_s.w", "f_s.weight").replace("f_s.b", "f_s.bias") k = k.replace("f.w", "f.weight").replace("f.b", "f.bias") k = k.replace("gn_1.g", "gn_1.weight").replace("gn_1.b", "gn_1.bias") k = k.replace("gn_2.g", "gn_2.weight").replace("gn_2.b", "gn_2.bias") k = k.replace("gn.g", "gn.weight").replace("gn.b", "gn.bias") return k def rename_state_dict(sd, embedding): sd = {rename_state_dict_key(k): v for k, v in sd.items()} sd["embed_time.emb.weight"] = embedding["weight"] return sd # encode with stable diffusion vae pipe = StableDiffusionPipeline.from_pretrained("runwayml/stable-diffusion-v1-5", torch_dtype=torch.float16) pipe.vae.cuda() # construct original decoder with jitted model decoder_consistency = ConsistencyDecoder(device="cuda:0") # construct UNet code, overwrite the decoder with conv_unet_vae model = ConvUNetVAE() model.load_state_dict( rename_state_dict( stl("consistency_decoder.safetensors"), stl("embedding.safetensors"), ) ) model = model.cuda() decoder_consistency.ckpt = model image = load_image(args.test_image, size=(256, 256), center_crop=True) latent = pipe.vae.encode(image.half().cuda()).latent_dist.sample() # decode with gan sample_gan = pipe.vae.decode(latent).sample.detach() save_image(sample_gan, "gan.png") # decode with conv_unet_vae sample_consistency_orig = decoder_consistency(latent, generator=torch.Generator("cpu").manual_seed(0)) save_image(sample_consistency_orig, "con_orig.png") ########### conversion print("CONVERSION") print("DOWN BLOCK ONE") block_one_sd_orig = model.down[0].state_dict() block_one_sd_new = {} for i in range(3): block_one_sd_new[f"resnets.{i}.norm1.weight"] = block_one_sd_orig.pop(f"{i}.gn_1.weight") block_one_sd_new[f"resnets.{i}.norm1.bias"] = block_one_sd_orig.pop(f"{i}.gn_1.bias") block_one_sd_new[f"resnets.{i}.conv1.weight"] = block_one_sd_orig.pop(f"{i}.f_1.weight") block_one_sd_new[f"resnets.{i}.conv1.bias"] = block_one_sd_orig.pop(f"{i}.f_1.bias") block_one_sd_new[f"resnets.{i}.time_emb_proj.weight"] = block_one_sd_orig.pop(f"{i}.f_t.weight") block_one_sd_new[f"resnets.{i}.time_emb_proj.bias"] = block_one_sd_orig.pop(f"{i}.f_t.bias") block_one_sd_new[f"resnets.{i}.norm2.weight"] = block_one_sd_orig.pop(f"{i}.gn_2.weight") block_one_sd_new[f"resnets.{i}.norm2.bias"] = block_one_sd_orig.pop(f"{i}.gn_2.bias") block_one_sd_new[f"resnets.{i}.conv2.weight"] = block_one_sd_orig.pop(f"{i}.f_2.weight") block_one_sd_new[f"resnets.{i}.conv2.bias"] = block_one_sd_orig.pop(f"{i}.f_2.bias") block_one_sd_new["downsamplers.0.norm1.weight"] = block_one_sd_orig.pop("3.gn_1.weight") block_one_sd_new["downsamplers.0.norm1.bias"] = block_one_sd_orig.pop("3.gn_1.bias") block_one_sd_new["downsamplers.0.conv1.weight"] = block_one_sd_orig.pop("3.f_1.weight") block_one_sd_new["downsamplers.0.conv1.bias"] = block_one_sd_orig.pop("3.f_1.bias") block_one_sd_new["downsamplers.0.time_emb_proj.weight"] = block_one_sd_orig.pop("3.f_t.weight") block_one_sd_new["downsamplers.0.time_emb_proj.bias"] = block_one_sd_orig.pop("3.f_t.bias") block_one_sd_new["downsamplers.0.norm2.weight"] = block_one_sd_orig.pop("3.gn_2.weight") block_one_sd_new["downsamplers.0.norm2.bias"] = block_one_sd_orig.pop("3.gn_2.bias") block_one_sd_new["downsamplers.0.conv2.weight"] = block_one_sd_orig.pop("3.f_2.weight") block_one_sd_new["downsamplers.0.conv2.bias"] = block_one_sd_orig.pop("3.f_2.bias") assert len(block_one_sd_orig) == 0 block_one = ResnetDownsampleBlock2D( in_channels=320, out_channels=320, temb_channels=1280, num_layers=3, add_downsample=True, resnet_time_scale_shift="scale_shift", resnet_eps=1e-5, ) block_one.load_state_dict(block_one_sd_new) print("DOWN BLOCK TWO") block_two_sd_orig = model.down[1].state_dict() block_two_sd_new = {} for i in range(3): block_two_sd_new[f"resnets.{i}.norm1.weight"] = block_two_sd_orig.pop(f"{i}.gn_1.weight") block_two_sd_new[f"resnets.{i}.norm1.bias"] = block_two_sd_orig.pop(f"{i}.gn_1.bias") block_two_sd_new[f"resnets.{i}.conv1.weight"] = block_two_sd_orig.pop(f"{i}.f_1.weight") block_two_sd_new[f"resnets.{i}.conv1.bias"] = block_two_sd_orig.pop(f"{i}.f_1.bias") block_two_sd_new[f"resnets.{i}.time_emb_proj.weight"] = block_two_sd_orig.pop(f"{i}.f_t.weight") block_two_sd_new[f"resnets.{i}.time_emb_proj.bias"] = block_two_sd_orig.pop(f"{i}.f_t.bias") block_two_sd_new[f"resnets.{i}.norm2.weight"] = block_two_sd_orig.pop(f"{i}.gn_2.weight") block_two_sd_new[f"resnets.{i}.norm2.bias"] = block_two_sd_orig.pop(f"{i}.gn_2.bias") block_two_sd_new[f"resnets.{i}.conv2.weight"] = block_two_sd_orig.pop(f"{i}.f_2.weight") block_two_sd_new[f"resnets.{i}.conv2.bias"] = block_two_sd_orig.pop(f"{i}.f_2.bias") if i == 0: block_two_sd_new[f"resnets.{i}.conv_shortcut.weight"] = block_two_sd_orig.pop(f"{i}.f_s.weight") block_two_sd_new[f"resnets.{i}.conv_shortcut.bias"] = block_two_sd_orig.pop(f"{i}.f_s.bias") block_two_sd_new["downsamplers.0.norm1.weight"] = block_two_sd_orig.pop("3.gn_1.weight") block_two_sd_new["downsamplers.0.norm1.bias"] = block_two_sd_orig.pop("3.gn_1.bias") block_two_sd_new["downsamplers.0.conv1.weight"] = block_two_sd_orig.pop("3.f_1.weight") block_two_sd_new["downsamplers.0.conv1.bias"] = block_two_sd_orig.pop("3.f_1.bias") block_two_sd_new["downsamplers.0.time_emb_proj.weight"] = block_two_sd_orig.pop("3.f_t.weight") block_two_sd_new["downsamplers.0.time_emb_proj.bias"] = block_two_sd_orig.pop("3.f_t.bias") block_two_sd_new["downsamplers.0.norm2.weight"] = block_two_sd_orig.pop("3.gn_2.weight") block_two_sd_new["downsamplers.0.norm2.bias"] = block_two_sd_orig.pop("3.gn_2.bias") block_two_sd_new["downsamplers.0.conv2.weight"] = block_two_sd_orig.pop("3.f_2.weight") block_two_sd_new["downsamplers.0.conv2.bias"] = block_two_sd_orig.pop("3.f_2.bias") assert len(block_two_sd_orig) == 0 block_two = ResnetDownsampleBlock2D( in_channels=320, out_channels=640, temb_channels=1280, num_layers=3, add_downsample=True, resnet_time_scale_shift="scale_shift", resnet_eps=1e-5, ) block_two.load_state_dict(block_two_sd_new) print("DOWN BLOCK THREE") block_three_sd_orig = model.down[2].state_dict() block_three_sd_new = {} for i in range(3): block_three_sd_new[f"resnets.{i}.norm1.weight"] = block_three_sd_orig.pop(f"{i}.gn_1.weight") block_three_sd_new[f"resnets.{i}.norm1.bias"] = block_three_sd_orig.pop(f"{i}.gn_1.bias") block_three_sd_new[f"resnets.{i}.conv1.weight"] = block_three_sd_orig.pop(f"{i}.f_1.weight") block_three_sd_new[f"resnets.{i}.conv1.bias"] = block_three_sd_orig.pop(f"{i}.f_1.bias") block_three_sd_new[f"resnets.{i}.time_emb_proj.weight"] = block_three_sd_orig.pop(f"{i}.f_t.weight") block_three_sd_new[f"resnets.{i}.time_emb_proj.bias"] = block_three_sd_orig.pop(f"{i}.f_t.bias") block_three_sd_new[f"resnets.{i}.norm2.weight"] = block_three_sd_orig.pop(f"{i}.gn_2.weight") block_three_sd_new[f"resnets.{i}.norm2.bias"] = block_three_sd_orig.pop(f"{i}.gn_2.bias") block_three_sd_new[f"resnets.{i}.conv2.weight"] = block_three_sd_orig.pop(f"{i}.f_2.weight") block_three_sd_new[f"resnets.{i}.conv2.bias"] = block_three_sd_orig.pop(f"{i}.f_2.bias") if i == 0: block_three_sd_new[f"resnets.{i}.conv_shortcut.weight"] = block_three_sd_orig.pop(f"{i}.f_s.weight") block_three_sd_new[f"resnets.{i}.conv_shortcut.bias"] = block_three_sd_orig.pop(f"{i}.f_s.bias") block_three_sd_new["downsamplers.0.norm1.weight"] = block_three_sd_orig.pop("3.gn_1.weight") block_three_sd_new["downsamplers.0.norm1.bias"] = block_three_sd_orig.pop("3.gn_1.bias") block_three_sd_new["downsamplers.0.conv1.weight"] = block_three_sd_orig.pop("3.f_1.weight") block_three_sd_new["downsamplers.0.conv1.bias"] = block_three_sd_orig.pop("3.f_1.bias") block_three_sd_new["downsamplers.0.time_emb_proj.weight"] = block_three_sd_orig.pop("3.f_t.weight") block_three_sd_new["downsamplers.0.time_emb_proj.bias"] = block_three_sd_orig.pop("3.f_t.bias") block_three_sd_new["downsamplers.0.norm2.weight"] = block_three_sd_orig.pop("3.gn_2.weight") block_three_sd_new["downsamplers.0.norm2.bias"] = block_three_sd_orig.pop("3.gn_2.bias") block_three_sd_new["downsamplers.0.conv2.weight"] = block_three_sd_orig.pop("3.f_2.weight") block_three_sd_new["downsamplers.0.conv2.bias"] = block_three_sd_orig.pop("3.f_2.bias") assert len(block_three_sd_orig) == 0 block_three = ResnetDownsampleBlock2D( in_channels=640, out_channels=1024, temb_channels=1280, num_layers=3, add_downsample=True, resnet_time_scale_shift="scale_shift", resnet_eps=1e-5, ) block_three.load_state_dict(block_three_sd_new) print("DOWN BLOCK FOUR") block_four_sd_orig = model.down[3].state_dict() block_four_sd_new = {} for i in range(3): block_four_sd_new[f"resnets.{i}.norm1.weight"] = block_four_sd_orig.pop(f"{i}.gn_1.weight") block_four_sd_new[f"resnets.{i}.norm1.bias"] = block_four_sd_orig.pop(f"{i}.gn_1.bias") block_four_sd_new[f"resnets.{i}.conv1.weight"] = block_four_sd_orig.pop(f"{i}.f_1.weight") block_four_sd_new[f"resnets.{i}.conv1.bias"] = block_four_sd_orig.pop(f"{i}.f_1.bias") block_four_sd_new[f"resnets.{i}.time_emb_proj.weight"] = block_four_sd_orig.pop(f"{i}.f_t.weight") block_four_sd_new[f"resnets.{i}.time_emb_proj.bias"] = block_four_sd_orig.pop(f"{i}.f_t.bias") block_four_sd_new[f"resnets.{i}.norm2.weight"] = block_four_sd_orig.pop(f"{i}.gn_2.weight") block_four_sd_new[f"resnets.{i}.norm2.bias"] = block_four_sd_orig.pop(f"{i}.gn_2.bias") block_four_sd_new[f"resnets.{i}.conv2.weight"] = block_four_sd_orig.pop(f"{i}.f_2.weight") block_four_sd_new[f"resnets.{i}.conv2.bias"] = block_four_sd_orig.pop(f"{i}.f_2.bias") assert len(block_four_sd_orig) == 0 block_four = ResnetDownsampleBlock2D( in_channels=1024, out_channels=1024, temb_channels=1280, num_layers=3, add_downsample=False, resnet_time_scale_shift="scale_shift", resnet_eps=1e-5, ) block_four.load_state_dict(block_four_sd_new) print("MID BLOCK 1") mid_block_one_sd_orig = model.mid.state_dict() mid_block_one_sd_new = {} for i in range(2): mid_block_one_sd_new[f"resnets.{i}.norm1.weight"] = mid_block_one_sd_orig.pop(f"{i}.gn_1.weight") mid_block_one_sd_new[f"resnets.{i}.norm1.bias"] = mid_block_one_sd_orig.pop(f"{i}.gn_1.bias") mid_block_one_sd_new[f"resnets.{i}.conv1.weight"] = mid_block_one_sd_orig.pop(f"{i}.f_1.weight") mid_block_one_sd_new[f"resnets.{i}.conv1.bias"] = mid_block_one_sd_orig.pop(f"{i}.f_1.bias") mid_block_one_sd_new[f"resnets.{i}.time_emb_proj.weight"] = mid_block_one_sd_orig.pop(f"{i}.f_t.weight") mid_block_one_sd_new[f"resnets.{i}.time_emb_proj.bias"] = mid_block_one_sd_orig.pop(f"{i}.f_t.bias") mid_block_one_sd_new[f"resnets.{i}.norm2.weight"] = mid_block_one_sd_orig.pop(f"{i}.gn_2.weight") mid_block_one_sd_new[f"resnets.{i}.norm2.bias"] = mid_block_one_sd_orig.pop(f"{i}.gn_2.bias") mid_block_one_sd_new[f"resnets.{i}.conv2.weight"] = mid_block_one_sd_orig.pop(f"{i}.f_2.weight") mid_block_one_sd_new[f"resnets.{i}.conv2.bias"] = mid_block_one_sd_orig.pop(f"{i}.f_2.bias") assert len(mid_block_one_sd_orig) == 0 mid_block_one = UNetMidBlock2D( in_channels=1024, temb_channels=1280, num_layers=1, resnet_time_scale_shift="scale_shift", resnet_eps=1e-5, add_attention=False, ) mid_block_one.load_state_dict(mid_block_one_sd_new) print("UP BLOCK ONE") up_block_one_sd_orig = model.up[-1].state_dict() up_block_one_sd_new = {} for i in range(4): up_block_one_sd_new[f"resnets.{i}.norm1.weight"] = up_block_one_sd_orig.pop(f"{i}.gn_1.weight") up_block_one_sd_new[f"resnets.{i}.norm1.bias"] = up_block_one_sd_orig.pop(f"{i}.gn_1.bias") up_block_one_sd_new[f"resnets.{i}.conv1.weight"] = up_block_one_sd_orig.pop(f"{i}.f_1.weight") up_block_one_sd_new[f"resnets.{i}.conv1.bias"] = up_block_one_sd_orig.pop(f"{i}.f_1.bias") up_block_one_sd_new[f"resnets.{i}.time_emb_proj.weight"] = up_block_one_sd_orig.pop(f"{i}.f_t.weight") up_block_one_sd_new[f"resnets.{i}.time_emb_proj.bias"] = up_block_one_sd_orig.pop(f"{i}.f_t.bias") up_block_one_sd_new[f"resnets.{i}.norm2.weight"] = up_block_one_sd_orig.pop(f"{i}.gn_2.weight") up_block_one_sd_new[f"resnets.{i}.norm2.bias"] = up_block_one_sd_orig.pop(f"{i}.gn_2.bias") up_block_one_sd_new[f"resnets.{i}.conv2.weight"] = up_block_one_sd_orig.pop(f"{i}.f_2.weight") up_block_one_sd_new[f"resnets.{i}.conv2.bias"] = up_block_one_sd_orig.pop(f"{i}.f_2.bias") up_block_one_sd_new[f"resnets.{i}.conv_shortcut.weight"] = up_block_one_sd_orig.pop(f"{i}.f_s.weight") up_block_one_sd_new[f"resnets.{i}.conv_shortcut.bias"] = up_block_one_sd_orig.pop(f"{i}.f_s.bias") up_block_one_sd_new["upsamplers.0.norm1.weight"] = up_block_one_sd_orig.pop("4.gn_1.weight") up_block_one_sd_new["upsamplers.0.norm1.bias"] = up_block_one_sd_orig.pop("4.gn_1.bias") up_block_one_sd_new["upsamplers.0.conv1.weight"] = up_block_one_sd_orig.pop("4.f_1.weight") up_block_one_sd_new["upsamplers.0.conv1.bias"] = up_block_one_sd_orig.pop("4.f_1.bias") up_block_one_sd_new["upsamplers.0.time_emb_proj.weight"] = up_block_one_sd_orig.pop("4.f_t.weight") up_block_one_sd_new["upsamplers.0.time_emb_proj.bias"] = up_block_one_sd_orig.pop("4.f_t.bias") up_block_one_sd_new["upsamplers.0.norm2.weight"] = up_block_one_sd_orig.pop("4.gn_2.weight") up_block_one_sd_new["upsamplers.0.norm2.bias"] = up_block_one_sd_orig.pop("4.gn_2.bias") up_block_one_sd_new["upsamplers.0.conv2.weight"] = up_block_one_sd_orig.pop("4.f_2.weight") up_block_one_sd_new["upsamplers.0.conv2.bias"] = up_block_one_sd_orig.pop("4.f_2.bias") assert len(up_block_one_sd_orig) == 0 up_block_one = ResnetUpsampleBlock2D( in_channels=1024, prev_output_channel=1024, out_channels=1024, temb_channels=1280, num_layers=4, add_upsample=True, resnet_time_scale_shift="scale_shift", resnet_eps=1e-5, ) up_block_one.load_state_dict(up_block_one_sd_new) print("UP BLOCK TWO") up_block_two_sd_orig = model.up[-2].state_dict() up_block_two_sd_new = {} for i in range(4): up_block_two_sd_new[f"resnets.{i}.norm1.weight"] = up_block_two_sd_orig.pop(f"{i}.gn_1.weight") up_block_two_sd_new[f"resnets.{i}.norm1.bias"] = up_block_two_sd_orig.pop(f"{i}.gn_1.bias") up_block_two_sd_new[f"resnets.{i}.conv1.weight"] = up_block_two_sd_orig.pop(f"{i}.f_1.weight") up_block_two_sd_new[f"resnets.{i}.conv1.bias"] = up_block_two_sd_orig.pop(f"{i}.f_1.bias") up_block_two_sd_new[f"resnets.{i}.time_emb_proj.weight"] = up_block_two_sd_orig.pop(f"{i}.f_t.weight") up_block_two_sd_new[f"resnets.{i}.time_emb_proj.bias"] = up_block_two_sd_orig.pop(f"{i}.f_t.bias") up_block_two_sd_new[f"resnets.{i}.norm2.weight"] = up_block_two_sd_orig.pop(f"{i}.gn_2.weight") up_block_two_sd_new[f"resnets.{i}.norm2.bias"] = up_block_two_sd_orig.pop(f"{i}.gn_2.bias") up_block_two_sd_new[f"resnets.{i}.conv2.weight"] = up_block_two_sd_orig.pop(f"{i}.f_2.weight") up_block_two_sd_new[f"resnets.{i}.conv2.bias"] = up_block_two_sd_orig.pop(f"{i}.f_2.bias") up_block_two_sd_new[f"resnets.{i}.conv_shortcut.weight"] = up_block_two_sd_orig.pop(f"{i}.f_s.weight") up_block_two_sd_new[f"resnets.{i}.conv_shortcut.bias"] = up_block_two_sd_orig.pop(f"{i}.f_s.bias") up_block_two_sd_new["upsamplers.0.norm1.weight"] = up_block_two_sd_orig.pop("4.gn_1.weight") up_block_two_sd_new["upsamplers.0.norm1.bias"] = up_block_two_sd_orig.pop("4.gn_1.bias") up_block_two_sd_new["upsamplers.0.conv1.weight"] = up_block_two_sd_orig.pop("4.f_1.weight") up_block_two_sd_new["upsamplers.0.conv1.bias"] = up_block_two_sd_orig.pop("4.f_1.bias") up_block_two_sd_new["upsamplers.0.time_emb_proj.weight"] = up_block_two_sd_orig.pop("4.f_t.weight") up_block_two_sd_new["upsamplers.0.time_emb_proj.bias"] = up_block_two_sd_orig.pop("4.f_t.bias") up_block_two_sd_new["upsamplers.0.norm2.weight"] = up_block_two_sd_orig.pop("4.gn_2.weight") up_block_two_sd_new["upsamplers.0.norm2.bias"] = up_block_two_sd_orig.pop("4.gn_2.bias") up_block_two_sd_new["upsamplers.0.conv2.weight"] = up_block_two_sd_orig.pop("4.f_2.weight") up_block_two_sd_new["upsamplers.0.conv2.bias"] = up_block_two_sd_orig.pop("4.f_2.bias") assert len(up_block_two_sd_orig) == 0 up_block_two = ResnetUpsampleBlock2D( in_channels=640, prev_output_channel=1024, out_channels=1024, temb_channels=1280, num_layers=4, add_upsample=True, resnet_time_scale_shift="scale_shift", resnet_eps=1e-5, ) up_block_two.load_state_dict(up_block_two_sd_new) print("UP BLOCK THREE") up_block_three_sd_orig = model.up[-3].state_dict() up_block_three_sd_new = {} for i in range(4): up_block_three_sd_new[f"resnets.{i}.norm1.weight"] = up_block_three_sd_orig.pop(f"{i}.gn_1.weight") up_block_three_sd_new[f"resnets.{i}.norm1.bias"] = up_block_three_sd_orig.pop(f"{i}.gn_1.bias") up_block_three_sd_new[f"resnets.{i}.conv1.weight"] = up_block_three_sd_orig.pop(f"{i}.f_1.weight") up_block_three_sd_new[f"resnets.{i}.conv1.bias"] = up_block_three_sd_orig.pop(f"{i}.f_1.bias") up_block_three_sd_new[f"resnets.{i}.time_emb_proj.weight"] = up_block_three_sd_orig.pop(f"{i}.f_t.weight") up_block_three_sd_new[f"resnets.{i}.time_emb_proj.bias"] = up_block_three_sd_orig.pop(f"{i}.f_t.bias") up_block_three_sd_new[f"resnets.{i}.norm2.weight"] = up_block_three_sd_orig.pop(f"{i}.gn_2.weight") up_block_three_sd_new[f"resnets.{i}.norm2.bias"] = up_block_three_sd_orig.pop(f"{i}.gn_2.bias") up_block_three_sd_new[f"resnets.{i}.conv2.weight"] = up_block_three_sd_orig.pop(f"{i}.f_2.weight") up_block_three_sd_new[f"resnets.{i}.conv2.bias"] = up_block_three_sd_orig.pop(f"{i}.f_2.bias") up_block_three_sd_new[f"resnets.{i}.conv_shortcut.weight"] = up_block_three_sd_orig.pop(f"{i}.f_s.weight") up_block_three_sd_new[f"resnets.{i}.conv_shortcut.bias"] = up_block_three_sd_orig.pop(f"{i}.f_s.bias") up_block_three_sd_new["upsamplers.0.norm1.weight"] = up_block_three_sd_orig.pop("4.gn_1.weight") up_block_three_sd_new["upsamplers.0.norm1.bias"] = up_block_three_sd_orig.pop("4.gn_1.bias") up_block_three_sd_new["upsamplers.0.conv1.weight"] = up_block_three_sd_orig.pop("4.f_1.weight") up_block_three_sd_new["upsamplers.0.conv1.bias"] = up_block_three_sd_orig.pop("4.f_1.bias") up_block_three_sd_new["upsamplers.0.time_emb_proj.weight"] = up_block_three_sd_orig.pop("4.f_t.weight") up_block_three_sd_new["upsamplers.0.time_emb_proj.bias"] = up_block_three_sd_orig.pop("4.f_t.bias") up_block_three_sd_new["upsamplers.0.norm2.weight"] = up_block_three_sd_orig.pop("4.gn_2.weight") up_block_three_sd_new["upsamplers.0.norm2.bias"] = up_block_three_sd_orig.pop("4.gn_2.bias") up_block_three_sd_new["upsamplers.0.conv2.weight"] = up_block_three_sd_orig.pop("4.f_2.weight") up_block_three_sd_new["upsamplers.0.conv2.bias"] = up_block_three_sd_orig.pop("4.f_2.bias") assert len(up_block_three_sd_orig) == 0 up_block_three = ResnetUpsampleBlock2D( in_channels=320, prev_output_channel=1024, out_channels=640, temb_channels=1280, num_layers=4, add_upsample=True, resnet_time_scale_shift="scale_shift", resnet_eps=1e-5, ) up_block_three.load_state_dict(up_block_three_sd_new) print("UP BLOCK FOUR") up_block_four_sd_orig = model.up[-4].state_dict() up_block_four_sd_new = {} for i in range(4): up_block_four_sd_new[f"resnets.{i}.norm1.weight"] = up_block_four_sd_orig.pop(f"{i}.gn_1.weight") up_block_four_sd_new[f"resnets.{i}.norm1.bias"] = up_block_four_sd_orig.pop(f"{i}.gn_1.bias") up_block_four_sd_new[f"resnets.{i}.conv1.weight"] = up_block_four_sd_orig.pop(f"{i}.f_1.weight") up_block_four_sd_new[f"resnets.{i}.conv1.bias"] = up_block_four_sd_orig.pop(f"{i}.f_1.bias") up_block_four_sd_new[f"resnets.{i}.time_emb_proj.weight"] = up_block_four_sd_orig.pop(f"{i}.f_t.weight") up_block_four_sd_new[f"resnets.{i}.time_emb_proj.bias"] = up_block_four_sd_orig.pop(f"{i}.f_t.bias") up_block_four_sd_new[f"resnets.{i}.norm2.weight"] = up_block_four_sd_orig.pop(f"{i}.gn_2.weight") up_block_four_sd_new[f"resnets.{i}.norm2.bias"] = up_block_four_sd_orig.pop(f"{i}.gn_2.bias") up_block_four_sd_new[f"resnets.{i}.conv2.weight"] = up_block_four_sd_orig.pop(f"{i}.f_2.weight") up_block_four_sd_new[f"resnets.{i}.conv2.bias"] = up_block_four_sd_orig.pop(f"{i}.f_2.bias") up_block_four_sd_new[f"resnets.{i}.conv_shortcut.weight"] = up_block_four_sd_orig.pop(f"{i}.f_s.weight") up_block_four_sd_new[f"resnets.{i}.conv_shortcut.bias"] = up_block_four_sd_orig.pop(f"{i}.f_s.bias") assert len(up_block_four_sd_orig) == 0 up_block_four = ResnetUpsampleBlock2D( in_channels=320, prev_output_channel=640, out_channels=320, temb_channels=1280, num_layers=4, add_upsample=False, resnet_time_scale_shift="scale_shift", resnet_eps=1e-5, ) up_block_four.load_state_dict(up_block_four_sd_new) print("initial projection (conv_in)") conv_in_sd_orig = model.embed_image.state_dict() conv_in_sd_new = {} conv_in_sd_new["weight"] = conv_in_sd_orig.pop("f.weight") conv_in_sd_new["bias"] = conv_in_sd_orig.pop("f.bias") assert len(conv_in_sd_orig) == 0 block_out_channels = [320, 640, 1024, 1024] in_channels = 7 conv_in_kernel = 3 conv_in_padding = (conv_in_kernel - 1) // 2 conv_in = nn.Conv2d(in_channels, block_out_channels[0], kernel_size=conv_in_kernel, padding=conv_in_padding) conv_in.load_state_dict(conv_in_sd_new) print("out projection (conv_out) (conv_norm_out)") out_channels = 6 norm_num_groups = 32 norm_eps = 1e-5 act_fn = "silu" conv_out_kernel = 3 conv_out_padding = (conv_out_kernel - 1) // 2 conv_norm_out = nn.GroupNorm(num_channels=block_out_channels[0], num_groups=norm_num_groups, eps=norm_eps) # uses torch.functional in orig # conv_act = get_activation(act_fn) conv_out = nn.Conv2d(block_out_channels[0], out_channels, kernel_size=conv_out_kernel, padding=conv_out_padding) conv_norm_out.load_state_dict(model.output.gn.state_dict()) conv_out.load_state_dict(model.output.f.state_dict()) print("timestep projection (time_proj) (time_embedding)") f1_sd = model.embed_time.f_1.state_dict() f2_sd = model.embed_time.f_2.state_dict() time_embedding_sd = { "linear_1.weight": f1_sd.pop("weight"), "linear_1.bias": f1_sd.pop("bias"), "linear_2.weight": f2_sd.pop("weight"), "linear_2.bias": f2_sd.pop("bias"), } assert len(f1_sd) == 0 assert len(f2_sd) == 0 time_embedding_type = "learned" num_train_timesteps = 1024 time_embedding_dim = 1280 time_proj = nn.Embedding(num_train_timesteps, block_out_channels[0]) timestep_input_dim = block_out_channels[0] time_embedding = TimestepEmbedding(timestep_input_dim, time_embedding_dim) time_proj.load_state_dict(model.embed_time.emb.state_dict()) time_embedding.load_state_dict(time_embedding_sd) print("CONVERT") time_embedding.to("cuda") time_proj.to("cuda") conv_in.to("cuda") block_one.to("cuda") block_two.to("cuda") block_three.to("cuda") block_four.to("cuda") mid_block_one.to("cuda") up_block_one.to("cuda") up_block_two.to("cuda") up_block_three.to("cuda") up_block_four.to("cuda") conv_norm_out.to("cuda") conv_out.to("cuda") model.time_proj = time_proj model.time_embedding = time_embedding model.embed_image = conv_in model.down[0] = block_one model.down[1] = block_two model.down[2] = block_three model.down[3] = block_four model.mid = mid_block_one model.up[-1] = up_block_one model.up[-2] = up_block_two model.up[-3] = up_block_three model.up[-4] = up_block_four model.output.gn = conv_norm_out model.output.f = conv_out model.converted = True sample_consistency_new = decoder_consistency(latent, generator=torch.Generator("cpu").manual_seed(0)) save_image(sample_consistency_new, "con_new.png") assert (sample_consistency_orig == sample_consistency_new).all() print("making unet") unet = UNet2DModel( in_channels=in_channels, out_channels=out_channels, down_block_types=( "ResnetDownsampleBlock2D", "ResnetDownsampleBlock2D", "ResnetDownsampleBlock2D", "ResnetDownsampleBlock2D", ), up_block_types=( "ResnetUpsampleBlock2D", "ResnetUpsampleBlock2D", "ResnetUpsampleBlock2D", "ResnetUpsampleBlock2D", ), block_out_channels=block_out_channels, layers_per_block=3, norm_num_groups=norm_num_groups, norm_eps=norm_eps, resnet_time_scale_shift="scale_shift", time_embedding_type="learned", num_train_timesteps=num_train_timesteps, add_attention=False, ) unet_state_dict = {} def add_state_dict(prefix, mod): for k, v in mod.state_dict().items(): unet_state_dict[f"{prefix}.{k}"] = v add_state_dict("conv_in", conv_in) add_state_dict("time_proj", time_proj) add_state_dict("time_embedding", time_embedding) add_state_dict("down_blocks.0", block_one) add_state_dict("down_blocks.1", block_two) add_state_dict("down_blocks.2", block_three) add_state_dict("down_blocks.3", block_four) add_state_dict("mid_block", mid_block_one) add_state_dict("up_blocks.0", up_block_one) add_state_dict("up_blocks.1", up_block_two) add_state_dict("up_blocks.2", up_block_three) add_state_dict("up_blocks.3", up_block_four) add_state_dict("conv_norm_out", conv_norm_out) add_state_dict("conv_out", conv_out) unet.load_state_dict(unet_state_dict) print("running with diffusers unet") unet.to("cuda") decoder_consistency.ckpt = unet sample_consistency_new_2 = decoder_consistency(latent, generator=torch.Generator("cpu").manual_seed(0)) save_image(sample_consistency_new_2, "con_new_2.png") assert (sample_consistency_orig == sample_consistency_new_2).all() print("running with diffusers model") Encoder.old_constructor = Encoder.__init__ def new_constructor(self, **kwargs): self.old_constructor(**kwargs) self.constructor_arguments = kwargs Encoder.__init__ = new_constructor vae = AutoencoderKL.from_pretrained("runwayml/stable-diffusion-v1-5", subfolder="vae") consistency_vae = ConsistencyDecoderVAE( encoder_args=vae.encoder.constructor_arguments, decoder_args=unet.config, scaling_factor=vae.config.scaling_factor, block_out_channels=vae.config.block_out_channels, latent_channels=vae.config.latent_channels, ) consistency_vae.encoder.load_state_dict(vae.encoder.state_dict()) consistency_vae.quant_conv.load_state_dict(vae.quant_conv.state_dict()) consistency_vae.decoder_unet.load_state_dict(unet.state_dict()) consistency_vae.to(dtype=torch.float16, device="cuda") sample_consistency_new_3 = consistency_vae.decode( 0.18215 * latent, generator=torch.Generator("cpu").manual_seed(0) ).sample print("max difference") print((sample_consistency_orig - sample_consistency_new_3).abs().max()) print("total difference") print((sample_consistency_orig - sample_consistency_new_3).abs().sum()) # assert (sample_consistency_orig == sample_consistency_new_3).all() print("running with diffusers pipeline") pipe = DiffusionPipeline.from_pretrained( "runwayml/stable-diffusion-v1-5", vae=consistency_vae, torch_dtype=torch.float16 ) pipe.to("cuda") pipe("horse", generator=torch.Generator("cpu").manual_seed(0)).images[0].save("horse.png") if args.save_pretrained is not None: consistency_vae.save_pretrained(args.save_pretrained)