import inspect

import os

from argparse import ArgumentParser

import numpy as np

import torch

from muse import MaskGiTUViT, VQGANModel

from muse import PipelineMuse as OldPipelineMuse

from transformers import CLIPTextModelWithProjection, CLIPTokenizer

from diffusers import VQModel

from diffusers.models.attention_processor import AttnProcessor

from diffusers.models.unets.uvit_2d import UVit2DModel

from diffusers.pipelines.amused.pipeline_amused import AmusedPipeline

from diffusers.schedulers import AmusedScheduler

torch.backends.cuda.enable_flash_sdp(False)

torch.backends.cuda.enable_mem_efficient_sdp(False)

torch.backends.cuda.enable_math_sdp(True)

os.environ["CUDA_LAUNCH_BLOCKING"] = "1"

os.environ["CUBLAS_WORKSPACE_CONFIG"] = ":16:8"

torch.use_deterministic_algorithms(True)

# Enable CUDNN deterministic mode

torch.backends.cudnn.deterministic = True

torch.backends.cudnn.benchmark = False

torch.backends.cuda.matmul.allow_tf32 = False

device = "cuda"

def main():

args = ArgumentParser()

args.add_argument("--model_256", action="store_true")

args.add_argument("--write_to", type=str, required=False, default=None)

args.add_argument("--transformer_path", type=str, required=False, default=None)

args = args.parse_args()

transformer_path = args.transformer_path

subfolder = "transformer"

if transformer_path is None:

if args.model_256:

transformer_path = "openMUSE/muse-256"

else:

transformer_path = (

"../research-run-512-checkpoints/research-run-512-with-downsample-checkpoint-554000/unwrapped_model/"

)

subfolder = None

old_transformer = MaskGiTUViT.from_pretrained(transformer_path, subfolder=subfolder)

old_transformer.to(device)

old_vae = VQGANModel.from_pretrained("openMUSE/muse-512", subfolder="vae")

old_vae.to(device)

vqvae = make_vqvae(old_vae)

tokenizer = CLIPTokenizer.from_pretrained("openMUSE/muse-512", subfolder="text_encoder")

text_encoder = CLIPTextModelWithProjection.from_pretrained("openMUSE/muse-512", subfolder="text_encoder")

text_encoder.to(device)

transformer = make_transformer(old_transformer, args.model_256)

scheduler = AmusedScheduler(mask_token_id=old_transformer.config.mask_token_id)

new_pipe = AmusedPipeline(

vqvae=vqvae, tokenizer=tokenizer, text_encoder=text_encoder, transformer=transformer, scheduler=scheduler

)

old_pipe = OldPipelineMuse(

vae=old_vae, transformer=old_transformer, text_encoder=text_encoder, tokenizer=tokenizer

)

old_pipe.to(device)

if args.model_256:

transformer_seq_len = 256

orig_size = (256, 256)

else:

transformer_seq_len = 1024

orig_size = (512, 512)

old_out = old_pipe(

"dog",

generator=torch.Generator(device).manual_seed(0),

transformer_seq_len=transformer_seq_len,

orig_size=orig_size,

timesteps=12,

)[0]

new_out = new_pipe("dog", generator=torch.Generator(device).manual_seed(0)).images[0]

old_out = np.array(old_out)

new_out = np.array(new_out)

diff = np.abs(old_out.astype(np.float64) - new_out.astype(np.float64))

# assert diff diff.sum() == 0

print("skipping pipeline full equivalence check")

print(f"max diff: {diff.max()}, diff.sum() / diff.size {diff.sum() / diff.size}")

if args.model_256:

assert diff.max() <= 3

assert diff.sum() / diff.size < 0.7

else:

assert diff.max() <= 1

assert diff.sum() / diff.size < 0.4

if args.write_to is not None:

new_pipe.save_pretrained(args.write_to)

def make_transformer(old_transformer, model_256):

args = dict(old_transformer.config)

force_down_up_sample = args["force_down_up_sample"]

signature = inspect.signature(UVit2DModel.__init__)

args_ = {

"downsample": force_down_up_sample,

"upsample": force_down_up_sample,

"block_out_channels": args["block_out_channels"][0],

"sample_size": 16 if model_256 else 32,

}

for s in list(signature.parameters.keys()):

if s in ["self", "downsample", "upsample", "sample_size", "block_out_channels"]:

continue

args_[s] = args[s]

new_transformer = UVit2DModel(**args_)

new_transformer.to(device)

new_transformer.set_attn_processor(AttnProcessor())

state_dict = old_transformer.state_dict()

state_dict["cond_embed.linear_1.weight"] = state_dict.pop("cond_embed.0.weight")

state_dict["cond_embed.linear_2.weight"] = state_dict.pop("cond_embed.2.weight")

for i in range(22):

state_dict[f"transformer_layers.{i}.norm1.norm.weight"] = state_dict.pop(

f"transformer_layers.{i}.attn_layer_norm.weight"

)

state_dict[f"transformer_layers.{i}.norm1.linear.weight"] = state_dict.pop(

f"transformer_layers.{i}.self_attn_adaLN_modulation.mapper.weight"

)

state_dict[f"transformer_layers.{i}.attn1.to_q.weight"] = state_dict.pop(

f"transformer_layers.{i}.attention.query.weight"

)

state_dict[f"transformer_layers.{i}.attn1.to_k.weight"] = state_dict.pop(

f"transformer_layers.{i}.attention.key.weight"

)

state_dict[f"transformer_layers.{i}.attn1.to_v.weight"] = state_dict.pop(

f"transformer_layers.{i}.attention.value.weight"

)

state_dict[f"transformer_layers.{i}.attn1.to_out.0.weight"] = state_dict.pop(

f"transformer_layers.{i}.attention.out.weight"

)

state_dict[f"transformer_layers.{i}.norm2.norm.weight"] = state_dict.pop(

f"transformer_layers.{i}.crossattn_layer_norm.weight"

)

state_dict[f"transformer_layers.{i}.norm2.linear.weight"] = state_dict.pop(

f"transformer_layers.{i}.cross_attn_adaLN_modulation.mapper.weight"

)

state_dict[f"transformer_layers.{i}.attn2.to_q.weight"] = state_dict.pop(

f"transformer_layers.{i}.crossattention.query.weight"

)

state_dict[f"transformer_layers.{i}.attn2.to_k.weight"] = state_dict.pop(

f"transformer_layers.{i}.crossattention.key.weight"

)

state_dict[f"transformer_layers.{i}.attn2.to_v.weight"] = state_dict.pop(

f"transformer_layers.{i}.crossattention.value.weight"

)

state_dict[f"transformer_layers.{i}.attn2.to_out.0.weight"] = state_dict.pop(

f"transformer_layers.{i}.crossattention.out.weight"

)

state_dict[f"transformer_layers.{i}.norm3.norm.weight"] = state_dict.pop(

f"transformer_layers.{i}.ffn.pre_mlp_layer_norm.weight"

)

state_dict[f"transformer_layers.{i}.norm3.linear.weight"] = state_dict.pop(

f"transformer_layers.{i}.ffn.adaLN_modulation.mapper.weight"

)

wi_0_weight = state_dict.pop(f"transformer_layers.{i}.ffn.wi_0.weight")

wi_1_weight = state_dict.pop(f"transformer_layers.{i}.ffn.wi_1.weight")

proj_weight = torch.concat([wi_1_weight, wi_0_weight], dim=0)

state_dict[f"transformer_layers.{i}.ff.net.0.proj.weight"] = proj_weight

state_dict[f"transformer_layers.{i}.ff.net.2.weight"] = state_dict.pop(f"transformer_layers.{i}.ffn.wo.weight")

if force_down_up_sample:

state_dict["down_block.downsample.norm.weight"] = state_dict.pop("down_blocks.0.downsample.0.norm.weight")

state_dict["down_block.downsample.conv.weight"] = state_dict.pop("down_blocks.0.downsample.1.weight")

state_dict["up_block.upsample.norm.weight"] = state_dict.pop("up_blocks.0.upsample.0.norm.weight")

state_dict["up_block.upsample.conv.weight"] = state_dict.pop("up_blocks.0.upsample.1.weight")

state_dict["mlm_layer.layer_norm.weight"] = state_dict.pop("mlm_layer.layer_norm.norm.weight")

for i in range(3):

state_dict[f"down_block.res_blocks.{i}.norm.weight"] = state_dict.pop(

f"down_blocks.0.res_blocks.{i}.norm.norm.weight"

)

state_dict[f"down_block.res_blocks.{i}.channelwise_linear_1.weight"] = state_dict.pop(

f"down_blocks.0.res_blocks.{i}.channelwise.0.weight"

)

state_dict[f"down_block.res_blocks.{i}.channelwise_norm.gamma"] = state_dict.pop(

f"down_blocks.0.res_blocks.{i}.channelwise.2.gamma"

)

state_dict[f"down_block.res_blocks.{i}.channelwise_norm.beta"] = state_dict.pop(

f"down_blocks.0.res_blocks.{i}.channelwise.2.beta"

)

state_dict[f"down_block.res_blocks.{i}.channelwise_linear_2.weight"] = state_dict.pop(

f"down_blocks.0.res_blocks.{i}.channelwise.4.weight"

)

state_dict[f"down_block.res_blocks.{i}.cond_embeds_mapper.weight"] = state_dict.pop(

f"down_blocks.0.res_blocks.{i}.adaLN_modulation.mapper.weight"

)

state_dict[f"down_block.attention_blocks.{i}.norm1.weight"] = state_dict.pop(

f"down_blocks.0.attention_blocks.{i}.attn_layer_norm.weight"

)

state_dict[f"down_block.attention_blocks.{i}.attn1.to_q.weight"] = state_dict.pop(

f"down_blocks.0.attention_blocks.{i}.attention.query.weight"

)

state_dict[f"down_block.attention_blocks.{i}.attn1.to_k.weight"] = state_dict.pop(

f"down_blocks.0.attention_blocks.{i}.attention.key.weight"

)

state_dict[f"down_block.attention_blocks.{i}.attn1.to_v.weight"] = state_dict.pop(

f"down_blocks.0.attention_blocks.{i}.attention.value.weight"

)

state_dict[f"down_block.attention_blocks.{i}.attn1.to_out.0.weight"] = state_dict.pop(

f"down_blocks.0.attention_blocks.{i}.attention.out.weight"

)

state_dict[f"down_block.attention_blocks.{i}.norm2.weight"] = state_dict.pop(

f"down_blocks.0.attention_blocks.{i}.crossattn_layer_norm.weight"

)

state_dict[f"down_block.attention_blocks.{i}.attn2.to_q.weight"] = state_dict.pop(

f"down_blocks.0.attention_blocks.{i}.crossattention.query.weight"

)

state_dict[f"down_block.attention_blocks.{i}.attn2.to_k.weight"] = state_dict.pop(

f"down_blocks.0.attention_blocks.{i}.crossattention.key.weight"

)

state_dict[f"down_block.attention_blocks.{i}.attn2.to_v.weight"] = state_dict.pop(

f"down_blocks.0.attention_blocks.{i}.crossattention.value.weight"

)

state_dict[f"down_block.attention_blocks.{i}.attn2.to_out.0.weight"] = state_dict.pop(

f"down_blocks.0.attention_blocks.{i}.crossattention.out.weight"

)

state_dict[f"up_block.res_blocks.{i}.norm.weight"] = state_dict.pop(

f"up_blocks.0.res_blocks.{i}.norm.norm.weight"

)

state_dict[f"up_block.res_blocks.{i}.channelwise_linear_1.weight"] = state_dict.pop(

f"up_blocks.0.res_blocks.{i}.channelwise.0.weight"

)

state_dict[f"up_block.res_blocks.{i}.channelwise_norm.gamma"] = state_dict.pop(

f"up_blocks.0.res_blocks.{i}.channelwise.2.gamma"

)

state_dict[f"up_block.res_blocks.{i}.channelwise_norm.beta"] = state_dict.pop(

f"up_blocks.0.res_blocks.{i}.channelwise.2.beta"

)

state_dict[f"up_block.res_blocks.{i}.channelwise_linear_2.weight"] = state_dict.pop(

f"up_blocks.0.res_blocks.{i}.channelwise.4.weight"

)

state_dict[f"up_block.res_blocks.{i}.cond_embeds_mapper.weight"] = state_dict.pop(

f"up_blocks.0.res_blocks.{i}.adaLN_modulation.mapper.weight"

)

state_dict[f"up_block.attention_blocks.{i}.norm1.weight"] = state_dict.pop(

f"up_blocks.0.attention_blocks.{i}.attn_layer_norm.weight"

)

state_dict[f"up_block.attention_blocks.{i}.attn1.to_q.weight"] = state_dict.pop(

f"up_blocks.0.attention_blocks.{i}.attention.query.weight"

)

state_dict[f"up_block.attention_blocks.{i}.attn1.to_k.weight"] = state_dict.pop(

f"up_blocks.0.attention_blocks.{i}.attention.key.weight"

)

state_dict[f"up_block.attention_blocks.{i}.attn1.to_v.weight"] = state_dict.pop(

f"up_blocks.0.attention_blocks.{i}.attention.value.weight"

)

state_dict[f"up_block.attention_blocks.{i}.attn1.to_out.0.weight"] = state_dict.pop(

f"up_blocks.0.attention_blocks.{i}.attention.out.weight"

)

state_dict[f"up_block.attention_blocks.{i}.norm2.weight"] = state_dict.pop(

f"up_blocks.0.attention_blocks.{i}.crossattn_layer_norm.weight"

)

state_dict[f"up_block.attention_blocks.{i}.attn2.to_q.weight"] = state_dict.pop(

f"up_blocks.0.attention_blocks.{i}.crossattention.query.weight"

)

state_dict[f"up_block.attention_blocks.{i}.attn2.to_k.weight"] = state_dict.pop(

f"up_blocks.0.attention_blocks.{i}.crossattention.key.weight"

)

state_dict[f"up_block.attention_blocks.{i}.attn2.to_v.weight"] = state_dict.pop(

f"up_blocks.0.attention_blocks.{i}.crossattention.value.weight"

)

state_dict[f"up_block.attention_blocks.{i}.attn2.to_out.0.weight"] = state_dict.pop(

f"up_blocks.0.attention_blocks.{i}.crossattention.out.weight"

)

for key in list(state_dict.keys()):

if key.startswith("up_blocks.0"):

key_ = "up_block." + ".".join(key.split(".")[2:])

state_dict[key_] = state_dict.pop(key)

if key.startswith("down_blocks.0"):

key_ = "down_block." + ".".join(key.split(".")[2:])

state_dict[key_] = state_dict.pop(key)

new_transformer.load_state_dict(state_dict)

input_ids = torch.randint(0, 10, (1, 32, 32), device=old_transformer.device)

encoder_hidden_states = torch.randn((1, 77, 768), device=old_transformer.device)

cond_embeds = torch.randn((1, 768), device=old_transformer.device)

micro_conds = torch.tensor([[512, 512, 0, 0, 6]], dtype=torch.float32, device=old_transformer.device)

old_out = old_transformer(input_ids.reshape(1, -1), encoder_hidden_states, cond_embeds, micro_conds)

old_out = old_out.reshape(1, 32, 32, 8192).permute(0, 3, 1, 2)

new_out = new_transformer(input_ids, encoder_hidden_states, cond_embeds, micro_conds)

# NOTE: these differences are solely due to using the geglu block that has a single linear layer of

# double output dimension instead of two different linear layers

max_diff = (old_out - new_out).abs().max()

total_diff = (old_out - new_out).abs().sum()

print(f"Transformer max_diff: {max_diff} total_diff: {total_diff}")

assert max_diff < 0.01

assert total_diff < 1500

return new_transformer

def make_vqvae(old_vae):

new_vae = VQModel(

act_fn="silu",

block_out_channels=[128, 256, 256, 512, 768],

down_block_types=[

"DownEncoderBlock2D",

"DownEncoderBlock2D",

"DownEncoderBlock2D",

"DownEncoderBlock2D",

"DownEncoderBlock2D",

],

in_channels=3,

latent_channels=64,

layers_per_block=2,

norm_num_groups=32,

num_vq_embeddings=8192,

out_channels=3,

sample_size=32,

up_block_types=[

"UpDecoderBlock2D",

"UpDecoderBlock2D",

"UpDecoderBlock2D",

"UpDecoderBlock2D",

"UpDecoderBlock2D",

],

mid_block_add_attention=False,

lookup_from_codebook=True,

)

new_vae.to(device)

# fmt: off

new_state_dict = {}

old_state_dict = old_vae.state_dict()

new_state_dict["encoder.conv_in.weight"] = old_state_dict.pop("encoder.conv_in.weight")

new_state_dict["encoder.conv_in.bias"] = old_state_dict.pop("encoder.conv_in.bias")

convert_vae_block_state_dict(old_state_dict, "encoder.down.0", new_state_dict, "encoder.down_blocks.0")

convert_vae_block_state_dict(old_state_dict, "encoder.down.1", new_state_dict, "encoder.down_blocks.1")

convert_vae_block_state_dict(old_state_dict, "encoder.down.2", new_state_dict, "encoder.down_blocks.2")

convert_vae_block_state_dict(old_state_dict, "encoder.down.3", new_state_dict, "encoder.down_blocks.3")

convert_vae_block_state_dict(old_state_dict, "encoder.down.4", new_state_dict, "encoder.down_blocks.4")

new_state_dict["encoder.mid_block.resnets.0.norm1.weight"] = old_state_dict.pop("encoder.mid.block_1.norm1.weight")

new_state_dict["encoder.mid_block.resnets.0.norm1.bias"] = old_state_dict.pop("encoder.mid.block_1.norm1.bias")

new_state_dict["encoder.mid_block.resnets.0.conv1.weight"] = old_state_dict.pop("encoder.mid.block_1.conv1.weight")

new_state_dict["encoder.mid_block.resnets.0.conv1.bias"] = old_state_dict.pop("encoder.mid.block_1.conv1.bias")

new_state_dict["encoder.mid_block.resnets.0.norm2.weight"] = old_state_dict.pop("encoder.mid.block_1.norm2.weight")

new_state_dict["encoder.mid_block.resnets.0.norm2.bias"] = old_state_dict.pop("encoder.mid.block_1.norm2.bias")

new_state_dict["encoder.mid_block.resnets.0.conv2.weight"] = old_state_dict.pop("encoder.mid.block_1.conv2.weight")

new_state_dict["encoder.mid_block.resnets.0.conv2.bias"] = old_state_dict.pop("encoder.mid.block_1.conv2.bias")

new_state_dict["encoder.mid_block.resnets.1.norm1.weight"] = old_state_dict.pop("encoder.mid.block_2.norm1.weight")

new_state_dict["encoder.mid_block.resnets.1.norm1.bias"] = old_state_dict.pop("encoder.mid.block_2.norm1.bias")

new_state_dict["encoder.mid_block.resnets.1.conv1.weight"] = old_state_dict.pop("encoder.mid.block_2.conv1.weight")

new_state_dict["encoder.mid_block.resnets.1.conv1.bias"] = old_state_dict.pop("encoder.mid.block_2.conv1.bias")

new_state_dict["encoder.mid_block.resnets.1.norm2.weight"] = old_state_dict.pop("encoder.mid.block_2.norm2.weight")

new_state_dict["encoder.mid_block.resnets.1.norm2.bias"] = old_state_dict.pop("encoder.mid.block_2.norm2.bias")

new_state_dict["encoder.mid_block.resnets.1.conv2.weight"] = old_state_dict.pop("encoder.mid.block_2.conv2.weight")

new_state_dict["encoder.mid_block.resnets.1.conv2.bias"] = old_state_dict.pop("encoder.mid.block_2.conv2.bias")

new_state_dict["encoder.conv_norm_out.weight"] = old_state_dict.pop("encoder.norm_out.weight")

new_state_dict["encoder.conv_norm_out.bias"] = old_state_dict.pop("encoder.norm_out.bias")

new_state_dict["encoder.conv_out.weight"] = old_state_dict.pop("encoder.conv_out.weight")

new_state_dict["encoder.conv_out.bias"] = old_state_dict.pop("encoder.conv_out.bias")

new_state_dict["quant_conv.weight"] = old_state_dict.pop("quant_conv.weight")

new_state_dict["quant_conv.bias"] = old_state_dict.pop("quant_conv.bias")

new_state_dict["quantize.embedding.weight"] = old_state_dict.pop("quantize.embedding.weight")

new_state_dict["post_quant_conv.weight"] = old_state_dict.pop("post_quant_conv.weight")

new_state_dict["post_quant_conv.bias"] = old_state_dict.pop("post_quant_conv.bias")

new_state_dict["decoder.conv_in.weight"] = old_state_dict.pop("decoder.conv_in.weight")

new_state_dict["decoder.conv_in.bias"] = old_state_dict.pop("decoder.conv_in.bias")

new_state_dict["decoder.mid_block.resnets.0.norm1.weight"] = old_state_dict.pop("decoder.mid.block_1.norm1.weight")

new_state_dict["decoder.mid_block.resnets.0.norm1.bias"] = old_state_dict.pop("decoder.mid.block_1.norm1.bias")

new_state_dict["decoder.mid_block.resnets.0.conv1.weight"] = old_state_dict.pop("decoder.mid.block_1.conv1.weight")

new_state_dict["decoder.mid_block.resnets.0.conv1.bias"] = old_state_dict.pop("decoder.mid.block_1.conv1.bias")

new_state_dict["decoder.mid_block.resnets.0.norm2.weight"] = old_state_dict.pop("decoder.mid.block_1.norm2.weight")

new_state_dict["decoder.mid_block.resnets.0.norm2.bias"] = old_state_dict.pop("decoder.mid.block_1.norm2.bias")

new_state_dict["decoder.mid_block.resnets.0.conv2.weight"] = old_state_dict.pop("decoder.mid.block_1.conv2.weight")

new_state_dict["decoder.mid_block.resnets.0.conv2.bias"] = old_state_dict.pop("decoder.mid.block_1.conv2.bias")

new_state_dict["decoder.mid_block.resnets.1.norm1.weight"] = old_state_dict.pop("decoder.mid.block_2.norm1.weight")

new_state_dict["decoder.mid_block.resnets.1.norm1.bias"] = old_state_dict.pop("decoder.mid.block_2.norm1.bias")

new_state_dict["decoder.mid_block.resnets.1.conv1.weight"] = old_state_dict.pop("decoder.mid.block_2.conv1.weight")

new_state_dict["decoder.mid_block.resnets.1.conv1.bias"] = old_state_dict.pop("decoder.mid.block_2.conv1.bias")

new_state_dict["decoder.mid_block.resnets.1.norm2.weight"] = old_state_dict.pop("decoder.mid.block_2.norm2.weight")

new_state_dict["decoder.mid_block.resnets.1.norm2.bias"] = old_state_dict.pop("decoder.mid.block_2.norm2.bias")

new_state_dict["decoder.mid_block.resnets.1.conv2.weight"] = old_state_dict.pop("decoder.mid.block_2.conv2.weight")

new_state_dict["decoder.mid_block.resnets.1.conv2.bias"] = old_state_dict.pop("decoder.mid.block_2.conv2.bias")

convert_vae_block_state_dict(old_state_dict, "decoder.up.0", new_state_dict, "decoder.up_blocks.4")

convert_vae_block_state_dict(old_state_dict, "decoder.up.1", new_state_dict, "decoder.up_blocks.3")

convert_vae_block_state_dict(old_state_dict, "decoder.up.2", new_state_dict, "decoder.up_blocks.2")

convert_vae_block_state_dict(old_state_dict, "decoder.up.3", new_state_dict, "decoder.up_blocks.1")

convert_vae_block_state_dict(old_state_dict, "decoder.up.4", new_state_dict, "decoder.up_blocks.0")

new_state_dict["decoder.conv_norm_out.weight"] = old_state_dict.pop("decoder.norm_out.weight")

new_state_dict["decoder.conv_norm_out.bias"] = old_state_dict.pop("decoder.norm_out.bias")

new_state_dict["decoder.conv_out.weight"] = old_state_dict.pop("decoder.conv_out.weight")

new_state_dict["decoder.conv_out.bias"] = old_state_dict.pop("decoder.conv_out.bias")

# fmt: on

assert len(old_state_dict.keys()) == 0

new_vae.load_state_dict(new_state_dict)

input = torch.randn((1, 3, 512, 512), device=device)

input = input.clamp(-1, 1)

old_encoder_output = old_vae.quant_conv(old_vae.encoder(input))

new_encoder_output = new_vae.quant_conv(new_vae.encoder(input))

assert (old_encoder_output == new_encoder_output).all()

old_decoder_output = old_vae.decoder(old_vae.post_quant_conv(old_encoder_output))

new_decoder_output = new_vae.decoder(new_vae.post_quant_conv(new_encoder_output))

# assert (old_decoder_output == new_decoder_output).all()

print("kipping vae decoder equivalence check")

print(f"vae decoder diff {(old_decoder_output - new_decoder_output).float().abs().sum()}")

old_output = old_vae(input)[0]

new_output = new_vae(input)[0]

# assert (old_output == new_output).all()

print("skipping full vae equivalence check")

print(f"vae full diff {(old_output - new_output).float().abs().sum()}")

return new_vae

def convert_vae_block_state_dict(old_state_dict, prefix_from, new_state_dict, prefix_to):

# fmt: off

new_state_dict[f"{prefix_to}.resnets.0.norm1.weight"] = old_state_dict.pop(f"{prefix_from}.block.0.norm1.weight")

new_state_dict[f"{prefix_to}.resnets.0.norm1.bias"] = old_state_dict.pop(f"{prefix_from}.block.0.norm1.bias")

new_state_dict[f"{prefix_to}.resnets.0.conv1.weight"] = old_state_dict.pop(f"{prefix_from}.block.0.conv1.weight")

new_state_dict[f"{prefix_to}.resnets.0.conv1.bias"] = old_state_dict.pop(f"{prefix_from}.block.0.conv1.bias")

new_state_dict[f"{prefix_to}.resnets.0.norm2.weight"] = old_state_dict.pop(f"{prefix_from}.block.0.norm2.weight")

new_state_dict[f"{prefix_to}.resnets.0.norm2.bias"] = old_state_dict.pop(f"{prefix_from}.block.0.norm2.bias")

new_state_dict[f"{prefix_to}.resnets.0.conv2.weight"] = old_state_dict.pop(f"{prefix_from}.block.0.conv2.weight")

new_state_dict[f"{prefix_to}.resnets.0.conv2.bias"] = old_state_dict.pop(f"{prefix_from}.block.0.conv2.bias")

if f"{prefix_from}.block.0.nin_shortcut.weight" in old_state_dict:

new_state_dict[f"{prefix_to}.resnets.0.conv_shortcut.weight"] = old_state_dict.pop(f"{prefix_from}.block.0.nin_shortcut.weight")

new_state_dict[f"{prefix_to}.resnets.0.conv_shortcut.bias"] = old_state_dict.pop(f"{prefix_from}.block.0.nin_shortcut.bias")

new_state_dict[f"{prefix_to}.resnets.1.norm1.weight"] = old_state_dict.pop(f"{prefix_from}.block.1.norm1.weight")

new_state_dict[f"{prefix_to}.resnets.1.norm1.bias"] = old_state_dict.pop(f"{prefix_from}.block.1.norm1.bias")

new_state_dict[f"{prefix_to}.resnets.1.conv1.weight"] = old_state_dict.pop(f"{prefix_from}.block.1.conv1.weight")

new_state_dict[f"{prefix_to}.resnets.1.conv1.bias"] = old_state_dict.pop(f"{prefix_from}.block.1.conv1.bias")

new_state_dict[f"{prefix_to}.resnets.1.norm2.weight"] = old_state_dict.pop(f"{prefix_from}.block.1.norm2.weight")

new_state_dict[f"{prefix_to}.resnets.1.norm2.bias"] = old_state_dict.pop(f"{prefix_from}.block.1.norm2.bias")

new_state_dict[f"{prefix_to}.resnets.1.conv2.weight"] = old_state_dict.pop(f"{prefix_from}.block.1.conv2.weight")

new_state_dict[f"{prefix_to}.resnets.1.conv2.bias"] = old_state_dict.pop(f"{prefix_from}.block.1.conv2.bias")

if f"{prefix_from}.downsample.conv.weight" in old_state_dict:

new_state_dict[f"{prefix_to}.downsamplers.0.conv.weight"] = old_state_dict.pop(f"{prefix_from}.downsample.conv.weight")

new_state_dict[f"{prefix_to}.downsamplers.0.conv.bias"] = old_state_dict.pop(f"{prefix_from}.downsample.conv.bias")

if f"{prefix_from}.upsample.conv.weight" in old_state_dict:

new_state_dict[f"{prefix_to}.upsamplers.0.conv.weight"] = old_state_dict.pop(f"{prefix_from}.upsample.conv.weight")

new_state_dict[f"{prefix_to}.upsamplers.0.conv.bias"] = old_state_dict.pop(f"{prefix_from}.upsample.conv.bias")

if f"{prefix_from}.block.2.norm1.weight" in old_state_dict:

new_state_dict[f"{prefix_to}.resnets.2.norm1.weight"] = old_state_dict.pop(f"{prefix_from}.block.2.norm1.weight")

new_state_dict[f"{prefix_to}.resnets.2.norm1.bias"] = old_state_dict.pop(f"{prefix_from}.block.2.norm1.bias")

new_state_dict[f"{prefix_to}.resnets.2.conv1.weight"] = old_state_dict.pop(f"{prefix_from}.block.2.conv1.weight")

new_state_dict[f"{prefix_to}.resnets.2.conv1.bias"] = old_state_dict.pop(f"{prefix_from}.block.2.conv1.bias")

new_state_dict[f"{prefix_to}.resnets.2.norm2.weight"] = old_state_dict.pop(f"{prefix_from}.block.2.norm2.weight")

new_state_dict[f"{prefix_to}.resnets.2.norm2.bias"] = old_state_dict.pop(f"{prefix_from}.block.2.norm2.bias")

new_state_dict[f"{prefix_to}.resnets.2.conv2.weight"] = old_state_dict.pop(f"{prefix_from}.block.2.conv2.weight")

new_state_dict[f"{prefix_to}.resnets.2.conv2.bias"] = old_state_dict.pop(f"{prefix_from}.block.2.conv2.bias")

# fmt: on

if __name__ == "__main__":

main()