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 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