Stable Diffusion V1.5 Inference With PyTorch Weights And More Features Like Stable Diffusion Web UI In TensorFlow
from PIL import Image

from stable_diffusion.stable_diffusion import StableDiffusion

model = StableDiffusion(img_height=512, img_width=512, jit_compile=True)
img = model.image_to_image(
    "a cute girl.",
    unconditional_guidance_scale=7.5,
    reference_image="/path/to/a_girl.jpg",
    reference_image_strength=0.8,
    num_steps=50,
)
Image.fromarray(img[0]).save("out.jpg")
print("Saved at out.jpg")

image_to_image.py
cpuimage/minSDTF
from PIL import Image

from stable_diffusion.stable_diffusion import StableDiffusion

model = StableDiffusion(img_height=512, img_width=512, jit_compile=True)
img = model.inpaint(
    "A dog with sunglasses, wearing comfy hat, looking at camera, highly detailed, ultra sharp, cinematic, 100mm lens, 8k resolution.",
    reference_image="/path/to/dog.jpg",
    inpaint_mask="/path/to/dog_mask.png",
    mask_blur_strength=5,
    unconditional_guidance_scale=8.0,
    reference_image_strength=0.9,
    num_steps=50,
)
Image.fromarray(img[0]).save("out.jpg")
print("Saved at out.jpg")

inpaint.py
cpuimage/minSDTF
# Copyright 2022 The KerasCV Authors
#
# Licensed under the Apache License, Version 2.0 (the "License");
# you may not use this file except in compliance with the License.
# You may obtain a copy of the License at
#
#     https://www.apache.org/licenses/LICENSE-2.0
#
# Unless required by applicable law or agreed to in writing, software
# distributed under the License is distributed on an "AS IS" BASIS,
# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
# See the License for the specific language governing permissions and
# limitations under the License.
from .clip_tokenizer import SimpleTokenizer
from .control_net import HintNet, ControlNet
from .diffusion_model import DiffusionModel
from .image_decoder import ImageDecoder
from .image_encoder import ImageEncoder
from .text_encoder import TextEncoder

stable_diffusion/__init__.py
cpuimage/minSDTF
# Copyright 2022 The KerasCV Authors
#
# Licensed under the Apache License, Version 2.0 (the "License");
# you may not use this file except in compliance with the License.
# You may obtain a copy of the License at
#
#     https://www.apache.org/licenses/LICENSE-2.0
#
# Unless required by applicable law or agreed to in writing, software
# distributed under the License is distributed on an "AS IS" BASIS,
# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
# See the License for the specific language governing permissions and
# limitations under the License.
import os

import numpy as np
import torch
from safetensors import safe_open

UNET_KEY_MAPPING = {
    'model.diffusion_model.time_embed.0.weight': 'time_embedding.linear_1.weight',
    'model.diffusion_model.time_embed.0.bias': 'time_embedding.linear_1.bias',
    'model.diffusion_model.input_blocks.0.0.weight': 'conv_in.weight',
    'model.diffusion_model.input_blocks.0.0.bias': 'conv_in.bias',
    'model.diffusion_model.time_embed.2.weight': 'time_embedding.linear_2.weight',
    'model.diffusion_model.time_embed.2.bias': 'time_embedding.linear_2.bias',
    'model.diffusion_model.input_blocks.1.0.in_layers.0.weight': 'down_blocks.0.resnets.0.norm1.weight',
    'model.diffusion_model.input_blocks.1.0.in_layers.0.bias': 'down_blocks.0.resnets.0.norm1.bias',
    'model.diffusion_model.input_blocks.1.0.in_layers.2.weight': 'down_blocks.0.resnets.0.conv1.weight',
    'model.diffusion_model.input_blocks.1.0.in_layers.2.bias': 'down_blocks.0.resnets.0.conv1.bias',
    'model.diffusion_model.input_blocks.1.0.emb_layers.1.weight': 'down_blocks.0.resnets.0.time_emb_proj.weight',
    'model.diffusion_model.input_blocks.1.0.emb_layers.1.bias': 'down_blocks.0.resnets.0.time_emb_proj.bias',
    'model.diffusion_model.input_blocks.1.0.out_layers.0.weight': 'down_blocks.0.resnets.0.norm2.weight',
    'model.diffusion_model.input_blocks.1.0.out_layers.0.bias': 'down_blocks.0.resnets.0.norm2.bias',
    'model.diffusion_model.input_blocks.1.0.out_layers.3.weight': 'down_blocks.0.resnets.0.conv2.weight',
    'model.diffusion_model.input_blocks.1.0.out_layers.3.bias': 'down_blocks.0.resnets.0.conv2.bias',
    'model.diffusion_model.input_blocks.1.1.norm.weight': 'down_blocks.0.attentions.0.norm.weight',
    'model.diffusion_model.input_blocks.1.1.norm.bias': 'down_blocks.0.attentions.0.norm.bias',
    'model.diffusion_model.input_blocks.1.1.proj_in.weight': 'down_blocks.0.attentions.0.proj_in.weight',
    'model.diffusion_model.input_blocks.1.1.proj_in.bias': 'down_blocks.0.attentions.0.proj_in.bias',
    'model.diffusion_model.input_blocks.1.1.transformer_blocks.0.norm1.weight': 'down_blocks.0.attentions.0.transformer_blocks.0.norm1.weight',
    'model.diffusion_model.input_blocks.1.1.transformer_blocks.0.norm1.bias': 'down_blocks.0.attentions.0.transformer_blocks.0.norm1.bias',
    'model.diffusion_model.input_blocks.1.1.transformer_blocks.0.attn1.to_q.weight': 'down_blocks.0.attentions.0.transformer_blocks.0.attn1.to_q.weight',
    'model.diffusion_model.input_blocks.1.1.transformer_blocks.0.attn1.to_k.weight': 'down_blocks.0.attentions.0.transformer_blocks.0.attn1.to_k.weight',
    'model.diffusion_model.input_blocks.1.1.transformer_blocks.0.attn1.to_v.weight': 'down_blocks.0.attentions.0.transformer_blocks.0.attn1.to_v.weight',
    'model.diffusion_model.input_blocks.1.1.transformer_blocks.0.attn1.to_out.0.weight': 'down_blocks.0.attentions.0.transformer_blocks.0.attn1.to_out.0.weight',
    'model.diffusion_model.input_blocks.1.1.transformer_blocks.0.attn1.to_out.0.bias': 'down_blocks.0.attentions.0.transformer_blocks.0.attn1.to_out.0.bias',
    'model.diffusion_model.input_blocks.1.1.transformer_blocks.0.norm2.weight': 'down_blocks.0.attentions.0.transformer_blocks.0.norm2.weight',
    'model.diffusion_model.input_blocks.1.1.transformer_blocks.0.norm2.bias': 'down_blocks.0.attentions.0.transformer_blocks.0.norm2.bias',
    'model.diffusion_model.input_blocks.1.1.transformer_blocks.0.attn2.to_q.weight': 'down_blocks.0.attentions.0.transformer_blocks.0.attn2.to_q.weight',
    'model.diffusion_model.input_blocks.1.1.transformer_blocks.0.attn2.to_k.weight': 'down_blocks.0.attentions.0.transformer_blocks.0.attn2.to_k.weight',
    'model.diffusion_model.input_blocks.1.1.transformer_blocks.0.attn2.to_v.weight': 'down_blocks.0.attentions.0.transformer_blocks.0.attn2.to_v.weight',
    'model.diffusion_model.input_blocks.1.1.transformer_blocks.0.attn2.to_out.0.weight': 'down_blocks.0.attentions.0.transformer_blocks.0.attn2.to_out.0.weight',
    'model.diffusion_model.input_blocks.1.1.transformer_blocks.0.attn2.to_out.0.bias': 'down_blocks.0.attentions.0.transformer_blocks.0.attn2.to_out.0.bias',
    'model.diffusion_model.input_blocks.1.1.transformer_blocks.0.norm3.weight': 'down_blocks.0.attentions.0.transformer_blocks.0.norm3.weight',
    'model.diffusion_model.input_blocks.1.1.transformer_blocks.0.norm3.bias': 'down_blocks.0.attentions.0.transformer_blocks.0.norm3.bias',
    'model.diffusion_model.input_blocks.1.1.transformer_blocks.0.ff.net.0.proj.weight': 'down_blocks.0.attentions.0.transformer_blocks.0.ff.net.0.proj.weight',
    'model.diffusion_model.input_blocks.1.1.transformer_blocks.0.ff.net.0.proj.bias': 'down_blocks.0.attentions.0.transformer_blocks.0.ff.net.0.proj.bias',
    'model.diffusion_model.input_blocks.1.1.transformer_blocks.0.ff.net.2.weight': 'down_blocks.0.attentions.0.transformer_blocks.0.ff.net.2.weight',
    'model.diffusion_model.input_blocks.1.1.transformer_blocks.0.ff.net.2.bias': 'down_blocks.0.attentions.0.transformer_blocks.0.ff.net.2.bias',
    'model.diffusion_model.input_blocks.1.1.proj_out.weight': 'down_blocks.0.attentions.0.proj_out.weight',
    'model.diffusion_model.input_blocks.1.1.proj_out.bias': 'down_blocks.0.attentions.0.proj_out.bias',
    'model.diffusion_model.input_blocks.2.0.in_layers.0.weight': 'down_blocks.0.resnets.1.norm1.weight',
    'model.diffusion_model.input_blocks.2.0.in_layers.0.bias': 'down_blocks.0.resnets.1.norm1.bias',
    'model.diffusion_model.input_blocks.2.0.in_layers.2.weight': 'down_blocks.0.resnets.1.conv1.weight',
    'model.diffusion_model.input_blocks.2.0.in_layers.2.bias': 'down_blocks.0.resnets.1.conv1.bias',
    'model.diffusion_model.input_blocks.2.0.emb_layers.1.weight': 'down_blocks.0.resnets.1.time_emb_proj.weight',
    'model.diffusion_model.input_blocks.2.0.emb_layers.1.bias': 'down_blocks.0.resnets.1.time_emb_proj.bias',
    'model.diffusion_model.input_blocks.2.0.out_layers.0.weight': 'down_blocks.0.resnets.1.norm2.weight',
    'model.diffusion_model.input_blocks.2.0.out_layers.0.bias': 'down_blocks.0.resnets.1.norm2.bias',
    'model.diffusion_model.input_blocks.2.0.out_layers.3.weight': 'down_blocks.0.resnets.1.conv2.weight',
    'model.diffusion_model.input_blocks.2.0.out_layers.3.bias': 'down_blocks.0.resnets.1.conv2.bias',
    'model.diffusion_model.input_blocks.2.1.norm.weight': 'down_blocks.0.attentions.1.norm.weight',
    'model.diffusion_model.input_blocks.2.1.norm.bias': 'down_blocks.0.attentions.1.norm.bias',
    'model.diffusion_model.input_blocks.2.1.proj_in.weight': 'down_blocks.0.attentions.1.proj_in.weight',
    'model.diffusion_model.input_blocks.2.1.proj_in.bias': 'down_blocks.0.attentions.1.proj_in.bias',
    'model.diffusion_model.input_blocks.2.1.transformer_blocks.0.norm1.weight': 'down_blocks.0.attentions.1.transformer_blocks.0.norm1.weight',
    'model.diffusion_model.input_blocks.2.1.transformer_blocks.0.norm1.bias': 'down_blocks.0.attentions.1.transformer_blocks.0.norm1.bias',
    'model.diffusion_model.input_blocks.2.1.transformer_blocks.0.attn1.to_q.weight': 'down_blocks.0.attentions.1.transformer_blocks.0.attn1.to_q.weight',
    'model.diffusion_model.input_blocks.2.1.transformer_blocks.0.attn1.to_k.weight': 'down_blocks.0.attentions.1.transformer_blocks.0.attn1.to_k.weight',
    'model.diffusion_model.input_blocks.2.1.transformer_blocks.0.attn1.to_v.weight': 'down_blocks.0.attentions.1.transformer_blocks.0.attn1.to_v.weight',
    'model.diffusion_model.input_blocks.2.1.transformer_blocks.0.attn1.to_out.0.weight': 'down_blocks.0.attentions.1.transformer_blocks.0.attn1.to_out.0.weight',
    'model.diffusion_model.input_blocks.2.1.transformer_blocks.0.attn1.to_out.0.bias': 'down_blocks.0.attentions.1.transformer_blocks.0.attn1.to_out.0.bias',
    'model.diffusion_model.input_blocks.2.1.transformer_blocks.0.norm2.weight': 'down_blocks.0.attentions.1.transformer_blocks.0.norm2.weight',
    'model.diffusion_model.input_blocks.2.1.transformer_blocks.0.norm2.bias': 'down_blocks.0.attentions.1.transformer_blocks.0.norm2.bias',
    'model.diffusion_model.input_blocks.2.1.transformer_blocks.0.attn2.to_q.weight': 'down_blocks.0.attentions.1.transformer_blocks.0.attn2.to_q.weight',
    'model.diffusion_model.input_blocks.2.1.transformer_blocks.0.attn2.to_k.weight': 'down_blocks.0.attentions.1.transformer_blocks.0.attn2.to_k.weight',
    'model.diffusion_model.input_blocks.2.1.transformer_blocks.0.attn2.to_v.weight': 'down_blocks.0.attentions.1.transformer_blocks.0.attn2.to_v.weight',
    'model.diffusion_model.input_blocks.2.1.transformer_blocks.0.attn2.to_out.0.weight': 'down_blocks.0.attentions.1.transformer_blocks.0.attn2.to_out.0.weight',
    'model.diffusion_model.input_blocks.2.1.transformer_blocks.0.attn2.to_out.0.bias': 'down_blocks.0.attentions.1.transformer_blocks.0.attn2.to_out.0.bias',
    'model.diffusion_model.input_blocks.2.1.transformer_blocks.0.norm3.weight': 'down_blocks.0.attentions.1.transformer_blocks.0.norm3.weight',
    'model.diffusion_model.input_blocks.2.1.transformer_blocks.0.norm3.bias': 'down_blocks.0.attentions.1.transformer_blocks.0.norm3.bias',
    'model.diffusion_model.input_blocks.2.1.transformer_blocks.0.ff.net.0.proj.weight': 'down_blocks.0.attentions.1.transformer_blocks.0.ff.net.0.proj.weight',
    'model.diffusion_model.input_blocks.2.1.transformer_blocks.0.ff.net.0.proj.bias': 'down_blocks.0.attentions.1.transformer_blocks.0.ff.net.0.proj.bias',
    'model.diffusion_model.input_blocks.2.1.transformer_blocks.0.ff.net.2.weight': 'down_blocks.0.attentions.1.transformer_blocks.0.ff.net.2.weight',
    'model.diffusion_model.input_blocks.2.1.transformer_blocks.0.ff.net.2.bias': 'down_blocks.0.attentions.1.transformer_blocks.0.ff.net.2.bias',
    'model.diffusion_model.input_blocks.2.1.proj_out.weight': 'down_blocks.0.attentions.1.proj_out.weight',
    'model.diffusion_model.input_blocks.2.1.proj_out.bias': 'down_blocks.0.attentions.1.proj_out.bias',
    'model.diffusion_model.input_blocks.3.0.op.weight': 'down_blocks.0.downsamplers.0.conv.weight',
    'model.diffusion_model.input_blocks.3.0.op.bias': 'down_blocks.0.downsamplers.0.conv.bias',
    'model.diffusion_model.input_blocks.4.0.in_layers.0.weight': 'down_blocks.1.resnets.0.norm1.weight',
    'model.diffusion_model.input_blocks.4.0.in_layers.0.bias': 'down_blocks.1.resnets.0.norm1.bias',
    'model.diffusion_model.input_blocks.4.0.in_layers.2.weight': 'down_blocks.1.resnets.0.conv1.weight',
    'model.diffusion_model.input_blocks.4.0.in_layers.2.bias': 'down_blocks.1.resnets.0.conv1.bias',
    'model.diffusion_model.input_blocks.4.0.emb_layers.1.weight': 'down_blocks.1.resnets.0.time_emb_proj.weight',
    'model.diffusion_model.input_blocks.4.0.emb_layers.1.bias': 'down_blocks.1.resnets.0.time_emb_proj.bias',
    'model.diffusion_model.input_blocks.4.0.out_layers.0.weight': 'down_blocks.1.resnets.0.norm2.weight',
    'model.diffusion_model.input_blocks.4.0.out_layers.0.bias': 'down_blocks.1.resnets.0.norm2.bias',
    'model.diffusion_model.input_blocks.4.0.out_layers.3.weight': 'down_blocks.1.resnets.0.conv2.weight',
    'model.diffusion_model.input_blocks.4.0.out_layers.3.bias': 'down_blocks.1.resnets.0.conv2.bias',
    'model.diffusion_model.input_blocks.4.0.skip_connection.weight': 'down_blocks.1.resnets.0.conv_shortcut.weight',
    'model.diffusion_model.input_blocks.4.0.skip_connection.bias': 'down_blocks.1.resnets.0.conv_shortcut.bias',
    'model.diffusion_model.input_blocks.4.1.norm.weight': 'down_blocks.1.attentions.0.norm.weight',
    'model.diffusion_model.input_blocks.4.1.norm.bias': 'down_blocks.1.attentions.0.norm.bias',
    'model.diffusion_model.input_blocks.4.1.proj_in.weight': 'down_blocks.1.attentions.0.proj_in.weight',
    'model.diffusion_model.input_blocks.4.1.proj_in.bias': 'down_blocks.1.attentions.0.proj_in.bias',
    'model.diffusion_model.input_blocks.4.1.transformer_blocks.0.norm1.weight': 'down_blocks.1.attentions.0.transformer_blocks.0.norm1.weight',
    'model.diffusion_model.input_blocks.4.1.transformer_blocks.0.norm1.bias': 'down_blocks.1.attentions.0.transformer_blocks.0.norm1.bias',
    'model.diffusion_model.input_blocks.4.1.transformer_blocks.0.attn1.to_q.weight': 'down_blocks.1.attentions.0.transformer_blocks.0.attn1.to_q.weight',
    'model.diffusion_model.input_blocks.4.1.transformer_blocks.0.attn1.to_k.weight': 'down_blocks.1.attentions.0.transformer_blocks.0.attn1.to_k.weight',
    'model.diffusion_model.input_blocks.4.1.transformer_blocks.0.attn1.to_v.weight': 'down_blocks.1.attentions.0.transformer_blocks.0.attn1.to_v.weight',
    'model.diffusion_model.input_blocks.4.1.transformer_blocks.0.attn1.to_out.0.weight': 'down_blocks.1.attentions.0.transformer_blocks.0.attn1.to_out.0.weight',
    'model.diffusion_model.input_blocks.4.1.transformer_blocks.0.attn1.to_out.0.bias': 'down_blocks.1.attentions.0.transformer_blocks.0.attn1.to_out.0.bias',
    'model.diffusion_model.input_blocks.4.1.transformer_blocks.0.norm2.weight': 'down_blocks.1.attentions.0.transformer_blocks.0.norm2.weight',
    'model.diffusion_model.input_blocks.4.1.transformer_blocks.0.norm2.bias': 'down_blocks.1.attentions.0.transformer_blocks.0.norm2.bias',
    'model.diffusion_model.input_blocks.4.1.transformer_blocks.0.attn2.to_q.weight': 'down_blocks.1.attentions.0.transformer_blocks.0.attn2.to_q.weight',
    'model.diffusion_model.input_blocks.4.1.transformer_blocks.0.attn2.to_k.weight': 'down_blocks.1.attentions.0.transformer_blocks.0.attn2.to_k.weight',
    'model.diffusion_model.input_blocks.4.1.transformer_blocks.0.attn2.to_v.weight': 'down_blocks.1.attentions.0.transformer_blocks.0.attn2.to_v.weight',
    'model.diffusion_model.input_blocks.4.1.transformer_blocks.0.attn2.to_out.0.weight': 'down_blocks.1.attentions.0.transformer_blocks.0.attn2.to_out.0.weight',
    'model.diffusion_model.input_blocks.4.1.transformer_blocks.0.attn2.to_out.0.bias': 'down_blocks.1.attentions.0.transformer_blocks.0.attn2.to_out.0.bias',
    'model.diffusion_model.input_blocks.4.1.transformer_blocks.0.norm3.weight': 'down_blocks.1.attentions.0.transformer_blocks.0.norm3.weight',
    'model.diffusion_model.input_blocks.4.1.transformer_blocks.0.norm3.bias': 'down_blocks.1.attentions.0.transformer_blocks.0.norm3.bias',
    'model.diffusion_model.input_blocks.4.1.transformer_blocks.0.ff.net.0.proj.weight': 'down_blocks.1.attentions.0.transformer_blocks.0.ff.net.0.proj.weight',
    'model.diffusion_model.input_blocks.4.1.transformer_blocks.0.ff.net.0.proj.bias': 'down_blocks.1.attentions.0.transformer_blocks.0.ff.net.0.proj.bias',
    'model.diffusion_model.input_blocks.4.1.transformer_blocks.0.ff.net.2.weight': 'down_blocks.1.attentions.0.transformer_blocks.0.ff.net.2.weight',
    'model.diffusion_model.input_blocks.4.1.transformer_blocks.0.ff.net.2.bias': 'down_blocks.1.attentions.0.transformer_blocks.0.ff.net.2.bias',
    'model.diffusion_model.input_blocks.4.1.proj_out.weight': 'down_blocks.1.attentions.0.proj_out.weight',
    'model.diffusion_model.input_blocks.4.1.proj_out.bias': 'down_blocks.1.attentions.0.proj_out.bias',
    'model.diffusion_model.input_blocks.5.0.in_layers.0.weight': 'down_blocks.1.resnets.1.norm1.weight',
    'model.diffusion_model.input_blocks.5.0.in_layers.0.bias': 'down_blocks.1.resnets.1.norm1.bias',
    'model.diffusion_model.input_blocks.5.0.in_layers.2.weight': 'down_blocks.1.resnets.1.conv1.weight',
    'model.diffusion_model.input_blocks.5.0.in_layers.2.bias': 'down_blocks.1.resnets.1.conv1.bias',
    'model.diffusion_model.input_blocks.5.0.emb_layers.1.weight': 'down_blocks.1.resnets.1.time_emb_proj.weight',
    'model.diffusion_model.input_blocks.5.0.emb_layers.1.bias': 'down_blocks.1.resnets.1.time_emb_proj.bias',
    'model.diffusion_model.input_blocks.5.0.out_layers.0.weight': 'down_blocks.1.resnets.1.norm2.weight',
    'model.diffusion_model.input_blocks.5.0.out_layers.0.bias': 'down_blocks.1.resnets.1.norm2.bias',
    'model.diffusion_model.input_blocks.5.0.out_layers.3.weight': 'down_blocks.1.resnets.1.conv2.weight',
    'model.diffusion_model.input_blocks.5.0.out_layers.3.bias': 'down_blocks.1.resnets.1.conv2.bias',
    'model.diffusion_model.input_blocks.5.1.norm.weight': 'down_blocks.1.attentions.1.norm.weight',
    'model.diffusion_model.input_blocks.5.1.norm.bias': 'down_blocks.1.attentions.1.norm.bias',
    'model.diffusion_model.input_blocks.5.1.proj_in.weight': 'down_blocks.1.attentions.1.proj_in.weight',
    'model.diffusion_model.input_blocks.5.1.proj_in.bias': 'down_blocks.1.attentions.1.proj_in.bias',
    'model.diffusion_model.input_blocks.5.1.transformer_blocks.0.norm1.weight': 'down_blocks.1.attentions.1.transformer_blocks.0.norm1.weight',
    'model.diffusion_model.input_blocks.5.1.transformer_blocks.0.norm1.bias': 'down_blocks.1.attentions.1.transformer_blocks.0.norm1.bias',
    'model.diffusion_model.input_blocks.5.1.transformer_blocks.0.attn1.to_q.weight': 'down_blocks.1.attentions.1.transformer_blocks.0.attn1.to_q.weight',
    'model.diffusion_model.input_blocks.5.1.transformer_blocks.0.attn1.to_k.weight': 'down_blocks.1.attentions.1.transformer_blocks.0.attn1.to_k.weight',
    'model.diffusion_model.input_blocks.5.1.transformer_blocks.0.attn1.to_v.weight': 'down_blocks.1.attentions.1.transformer_blocks.0.attn1.to_v.weight',
    'model.diffusion_model.input_blocks.5.1.transformer_blocks.0.attn1.to_out.0.weight': 'down_blocks.1.attentions.1.transformer_blocks.0.attn1.to_out.0.weight',
    'model.diffusion_model.input_blocks.5.1.transformer_blocks.0.attn1.to_out.0.bias': 'down_blocks.1.attentions.1.transformer_blocks.0.attn1.to_out.0.bias',
    'model.diffusion_model.input_blocks.5.1.transformer_blocks.0.norm2.weight': 'down_blocks.1.attentions.1.transformer_blocks.0.norm2.weight',
    'model.diffusion_model.input_blocks.5.1.transformer_blocks.0.norm2.bias': 'down_blocks.1.attentions.1.transformer_blocks.0.norm2.bias',
    'model.diffusion_model.input_blocks.5.1.transformer_blocks.0.attn2.to_q.weight': 'down_blocks.1.attentions.1.transformer_blocks.0.attn2.to_q.weight',
    'model.diffusion_model.input_blocks.5.1.transformer_blocks.0.attn2.to_k.weight': 'down_blocks.1.attentions.1.transformer_blocks.0.attn2.to_k.weight',
    'model.diffusion_model.input_blocks.5.1.transformer_blocks.0.attn2.to_v.weight': 'down_blocks.1.attentions.1.transformer_blocks.0.attn2.to_v.weight',
    'model.diffusion_model.input_blocks.5.1.transformer_blocks.0.attn2.to_out.0.weight': 'down_blocks.1.attentions.1.transformer_blocks.0.attn2.to_out.0.weight',
    'model.diffusion_model.input_blocks.5.1.transformer_blocks.0.attn2.to_out.0.bias': 'down_blocks.1.attentions.1.transformer_blocks.0.attn2.to_out.0.bias',
    'model.diffusion_model.input_blocks.5.1.transformer_blocks.0.norm3.weight': 'down_blocks.1.attentions.1.transformer_blocks.0.norm3.weight',
    'model.diffusion_model.input_blocks.5.1.transformer_blocks.0.norm3.bias': 'down_blocks.1.attentions.1.transformer_blocks.0.norm3.bias',
    'model.diffusion_model.input_blocks.5.1.transformer_blocks.0.ff.net.0.proj.weight': 'down_blocks.1.attentions.1.transformer_blocks.0.ff.net.0.proj.weight',
    'model.diffusion_model.input_blocks.5.1.transformer_blocks.0.ff.net.0.proj.bias': 'down_blocks.1.attentions.1.transformer_blocks.0.ff.net.0.proj.bias',
    'model.diffusion_model.input_blocks.5.1.transformer_blocks.0.ff.net.2.weight': 'down_blocks.1.attentions.1.transformer_blocks.0.ff.net.2.weight',
    'model.diffusion_model.input_blocks.5.1.transformer_blocks.0.ff.net.2.bias': 'down_blocks.1.attentions.1.transformer_blocks.0.ff.net.2.bias',
    'model.diffusion_model.input_blocks.5.1.proj_out.weight': 'down_blocks.1.attentions.1.proj_out.weight',
    'model.diffusion_model.input_blocks.5.1.proj_out.bias': 'down_blocks.1.attentions.1.proj_out.bias',
    'model.diffusion_model.input_blocks.6.0.op.weight': 'down_blocks.1.downsamplers.0.conv.weight',
    'model.diffusion_model.input_blocks.6.0.op.bias': 'down_blocks.1.downsamplers.0.conv.bias',
    'model.diffusion_model.input_blocks.7.0.in_layers.0.weight': 'down_blocks.2.resnets.0.norm1.weight',
    'model.diffusion_model.input_blocks.7.0.in_layers.0.bias': 'down_blocks.2.resnets.0.norm1.bias',
    'model.diffusion_model.input_blocks.7.0.in_layers.2.weight': 'down_blocks.2.resnets.0.conv1.weight',
    'model.diffusion_model.input_blocks.7.0.in_layers.2.bias': 'down_blocks.2.resnets.0.conv1.bias',
    'model.diffusion_model.input_blocks.7.0.emb_layers.1.weight': 'down_blocks.2.resnets.0.time_emb_proj.weight',
    'model.diffusion_model.input_blocks.7.0.emb_layers.1.bias': 'down_blocks.2.resnets.0.time_emb_proj.bias',
    'model.diffusion_model.input_blocks.7.0.out_layers.0.weight': 'down_blocks.2.resnets.0.norm2.weight',
    'model.diffusion_model.input_blocks.7.0.out_layers.0.bias': 'down_blocks.2.resnets.0.norm2.bias',
    'model.diffusion_model.input_blocks.7.0.out_layers.3.weight': 'down_blocks.2.resnets.0.conv2.weight',
    'model.diffusion_model.input_blocks.7.0.out_layers.3.bias': 'down_blocks.2.resnets.0.conv2.bias',
    'model.diffusion_model.input_blocks.7.0.skip_connection.weight': 'down_blocks.2.resnets.0.conv_shortcut.weight',
    'model.diffusion_model.input_blocks.7.0.skip_connection.bias': 'down_blocks.2.resnets.0.conv_shortcut.bias',
    'model.diffusion_model.input_blocks.7.1.norm.weight': 'down_blocks.2.attentions.0.norm.weight',
    'model.diffusion_model.input_blocks.7.1.norm.bias': 'down_blocks.2.attentions.0.norm.bias',
    'model.diffusion_model.input_blocks.7.1.proj_in.weight': 'down_blocks.2.attentions.0.proj_in.weight',
    'model.diffusion_model.input_blocks.7.1.proj_in.bias': 'down_blocks.2.attentions.0.proj_in.bias',
    'model.diffusion_model.input_blocks.7.1.transformer_blocks.0.norm1.weight': 'down_blocks.2.attentions.0.transformer_blocks.0.norm1.weight',
    'model.diffusion_model.input_blocks.7.1.transformer_blocks.0.norm1.bias': 'down_blocks.2.attentions.0.transformer_blocks.0.norm1.bias',
    'model.diffusion_model.input_blocks.7.1.transformer_blocks.0.attn1.to_q.weight': 'down_blocks.2.attentions.0.transformer_blocks.0.attn1.to_q.weight',
    'model.diffusion_model.input_blocks.7.1.transformer_blocks.0.attn1.to_k.weight': 'down_blocks.2.attentions.0.transformer_blocks.0.attn1.to_k.weight',
    'model.diffusion_model.input_blocks.7.1.transformer_blocks.0.attn1.to_v.weight': 'down_blocks.2.attentions.0.transformer_blocks.0.attn1.to_v.weight',
    'model.diffusion_model.input_blocks.7.1.transformer_blocks.0.attn1.to_out.0.weight': 'down_blocks.2.attentions.0.transformer_blocks.0.attn1.to_out.0.weight',
    'model.diffusion_model.input_blocks.7.1.transformer_blocks.0.attn1.to_out.0.bias': 'down_blocks.2.attentions.0.transformer_blocks.0.attn1.to_out.0.bias',
    'model.diffusion_model.input_blocks.7.1.transformer_blocks.0.norm2.weight': 'down_blocks.2.attentions.0.transformer_blocks.0.norm2.weight',
    'model.diffusion_model.input_blocks.7.1.transformer_blocks.0.norm2.bias': 'down_blocks.2.attentions.0.transformer_blocks.0.norm2.bias',
    'model.diffusion_model.input_blocks.7.1.transformer_blocks.0.attn2.to_q.weight': 'down_blocks.2.attentions.0.transformer_blocks.0.attn2.to_q.weight',
    'model.diffusion_model.input_blocks.7.1.transformer_blocks.0.attn2.to_k.weight': 'down_blocks.2.attentions.0.transformer_blocks.0.attn2.to_k.weight',
    'model.diffusion_model.input_blocks.7.1.transformer_blocks.0.attn2.to_v.weight': 'down_blocks.2.attentions.0.transformer_blocks.0.attn2.to_v.weight',
    'model.diffusion_model.input_blocks.7.1.transformer_blocks.0.attn2.to_out.0.weight': 'down_blocks.2.attentions.0.transformer_blocks.0.attn2.to_out.0.weight',
    'model.diffusion_model.input_blocks.7.1.transformer_blocks.0.attn2.to_out.0.bias': 'down_blocks.2.attentions.0.transformer_blocks.0.attn2.to_out.0.bias',
    'model.diffusion_model.input_blocks.7.1.transformer_blocks.0.norm3.weight': 'down_blocks.2.attentions.0.transformer_blocks.0.norm3.weight',
    'model.diffusion_model.input_blocks.7.1.transformer_blocks.0.norm3.bias': 'down_blocks.2.attentions.0.transformer_blocks.0.norm3.bias',
    'model.diffusion_model.input_blocks.7.1.transformer_blocks.0.ff.net.0.proj.weight': 'down_blocks.2.attentions.0.transformer_blocks.0.ff.net.0.proj.weight',
    'model.diffusion_model.input_blocks.7.1.transformer_blocks.0.ff.net.0.proj.bias': 'down_blocks.2.attentions.0.transformer_blocks.0.ff.net.0.proj.bias',
    'model.diffusion_model.input_blocks.7.1.transformer_blocks.0.ff.net.2.weight': 'down_blocks.2.attentions.0.transformer_blocks.0.ff.net.2.weight',
    'model.diffusion_model.input_blocks.7.1.transformer_blocks.0.ff.net.2.bias': 'down_blocks.2.attentions.0.transformer_blocks.0.ff.net.2.bias',
    'model.diffusion_model.input_blocks.7.1.proj_out.weight': 'down_blocks.2.attentions.0.proj_out.weight',
    'model.diffusion_model.input_blocks.7.1.proj_out.bias': 'down_blocks.2.attentions.0.proj_out.bias',
    'model.diffusion_model.input_blocks.8.0.in_layers.0.weight': 'down_blocks.2.resnets.1.norm1.weight',
    'model.diffusion_model.input_blocks.8.0.in_layers.0.bias': 'down_blocks.2.resnets.1.norm1.bias',
    'model.diffusion_model.input_blocks.8.0.in_layers.2.weight': 'down_blocks.2.resnets.1.conv1.weight',
    'model.diffusion_model.input_blocks.8.0.in_layers.2.bias': 'down_blocks.2.resnets.1.conv1.bias',
    'model.diffusion_model.input_blocks.8.0.emb_layers.1.weight': 'down_blocks.2.resnets.1.time_emb_proj.weight',
    'model.diffusion_model.input_blocks.8.0.emb_layers.1.bias': 'down_blocks.2.resnets.1.time_emb_proj.bias',
    'model.diffusion_model.input_blocks.8.0.out_layers.0.weight': 'down_blocks.2.resnets.1.norm2.weight',
    'model.diffusion_model.input_blocks.8.0.out_layers.0.bias': 'down_blocks.2.resnets.1.norm2.bias',
    'model.diffusion_model.input_blocks.8.0.out_layers.3.weight': 'down_blocks.2.resnets.1.conv2.weight',
    'model.diffusion_model.input_blocks.8.0.out_layers.3.bias': 'down_blocks.2.resnets.1.conv2.bias',
    'model.diffusion_model.input_blocks.8.1.norm.weight': 'down_blocks.2.attentions.1.norm.weight',
    'model.diffusion_model.input_blocks.8.1.norm.bias': 'down_blocks.2.attentions.1.norm.bias',
    'model.diffusion_model.input_blocks.8.1.proj_in.weight': 'down_blocks.2.attentions.1.proj_in.weight',
    'model.diffusion_model.input_blocks.8.1.proj_in.bias': 'down_blocks.2.attentions.1.proj_in.bias',
    'model.diffusion_model.input_blocks.8.1.transformer_blocks.0.norm1.weight': 'down_blocks.2.attentions.1.transformer_blocks.0.norm1.weight',
    'model.diffusion_model.input_blocks.8.1.transformer_blocks.0.norm1.bias': 'down_blocks.2.attentions.1.transformer_blocks.0.norm1.bias',
    'model.diffusion_model.input_blocks.8.1.transformer_blocks.0.attn1.to_q.weight': 'down_blocks.2.attentions.1.transformer_blocks.0.attn1.to_q.weight',
    'model.diffusion_model.input_blocks.8.1.transformer_blocks.0.attn1.to_k.weight': 'down_blocks.2.attentions.1.transformer_blocks.0.attn1.to_k.weight',
    'model.diffusion_model.input_blocks.8.1.transformer_blocks.0.attn1.to_v.weight': 'down_blocks.2.attentions.1.transformer_blocks.0.attn1.to_v.weight',
    'model.diffusion_model.input_blocks.8.1.transformer_blocks.0.attn1.to_out.0.weight': 'down_blocks.2.attentions.1.transformer_blocks.0.attn1.to_out.0.weight',
    'model.diffusion_model.input_blocks.8.1.transformer_blocks.0.attn1.to_out.0.bias': 'down_blocks.2.attentions.1.transformer_blocks.0.attn1.to_out.0.bias',
    'model.diffusion_model.input_blocks.8.1.transformer_blocks.0.norm2.weight': 'down_blocks.2.attentions.1.transformer_blocks.0.norm2.weight',
    'model.diffusion_model.input_blocks.8.1.transformer_blocks.0.norm2.bias': 'down_blocks.2.attentions.1.transformer_blocks.0.norm2.bias',
    'model.diffusion_model.input_blocks.8.1.transformer_blocks.0.attn2.to_q.weight': 'down_blocks.2.attentions.1.transformer_blocks.0.attn2.to_q.weight',
    'model.diffusion_model.input_blocks.8.1.transformer_blocks.0.attn2.to_k.weight': 'down_blocks.2.attentions.1.transformer_blocks.0.attn2.to_k.weight',
    'model.diffusion_model.input_blocks.8.1.transformer_blocks.0.attn2.to_v.weight': 'down_blocks.2.attentions.1.transformer_blocks.0.attn2.to_v.weight',
    'model.diffusion_model.input_blocks.8.1.transformer_blocks.0.attn2.to_out.0.weight': 'down_blocks.2.attentions.1.transformer_blocks.0.attn2.to_out.0.weight',
    'model.diffusion_model.input_blocks.8.1.transformer_blocks.0.attn2.to_out.0.bias': 'down_blocks.2.attentions.1.transformer_blocks.0.attn2.to_out.0.bias',
    'model.diffusion_model.input_blocks.8.1.transformer_blocks.0.norm3.weight': 'down_blocks.2.attentions.1.transformer_blocks.0.norm3.weight',
    'model.diffusion_model.input_blocks.8.1.transformer_blocks.0.norm3.bias': 'down_blocks.2.attentions.1.transformer_blocks.0.norm3.bias',
    'model.diffusion_model.input_blocks.8.1.transformer_blocks.0.ff.net.0.proj.weight': 'down_blocks.2.attentions.1.transformer_blocks.0.ff.net.0.proj.weight',
    'model.diffusion_model.input_blocks.8.1.transformer_blocks.0.ff.net.0.proj.bias': 'down_blocks.2.attentions.1.transformer_blocks.0.ff.net.0.proj.bias',
    'model.diffusion_model.input_blocks.8.1.transformer_blocks.0.ff.net.2.weight': 'down_blocks.2.attentions.1.transformer_blocks.0.ff.net.2.weight',
    'model.diffusion_model.input_blocks.8.1.transformer_blocks.0.ff.net.2.bias': 'down_blocks.2.attentions.1.transformer_blocks.0.ff.net.2.bias',
    'model.diffusion_model.input_blocks.8.1.proj_out.weight': 'down_blocks.2.attentions.1.proj_out.weight',
    'model.diffusion_model.input_blocks.8.1.proj_out.bias': 'down_blocks.2.attentions.1.proj_out.bias',
    'model.diffusion_model.input_blocks.9.0.op.weight': 'down_blocks.2.downsamplers.0.conv.weight',
    'model.diffusion_model.input_blocks.9.0.op.bias': 'down_blocks.2.downsamplers.0.conv.bias',
    'model.diffusion_model.input_blocks.10.0.in_layers.0.weight': 'down_blocks.3.resnets.0.norm1.weight',
    'model.diffusion_model.input_blocks.10.0.in_layers.0.bias': 'down_blocks.3.resnets.0.norm1.bias',
    'model.diffusion_model.input_blocks.10.0.in_layers.2.weight': 'down_blocks.3.resnets.0.conv1.weight',
    'model.diffusion_model.input_blocks.10.0.in_layers.2.bias': 'down_blocks.3.resnets.0.conv1.bias',
    'model.diffusion_model.input_blocks.10.0.emb_layers.1.weight': 'down_blocks.3.resnets.0.time_emb_proj.weight',
    'model.diffusion_model.input_blocks.10.0.emb_layers.1.bias': 'down_blocks.3.resnets.0.time_emb_proj.bias',
    'model.diffusion_model.input_blocks.10.0.out_layers.0.weight': 'down_blocks.3.resnets.0.norm2.weight',
    'model.diffusion_model.input_blocks.10.0.out_layers.0.bias': 'down_blocks.3.resnets.0.norm2.bias',
    'model.diffusion_model.input_blocks.10.0.out_layers.3.weight': 'down_blocks.3.resnets.0.conv2.weight',
    'model.diffusion_model.input_blocks.10.0.out_layers.3.bias': 'down_blocks.3.resnets.0.conv2.bias',
    'model.diffusion_model.input_blocks.11.0.in_layers.0.weight': 'down_blocks.3.resnets.1.norm1.weight',
    'model.diffusion_model.input_blocks.11.0.in_layers.0.bias': 'down_blocks.3.resnets.1.norm1.bias',
    'model.diffusion_model.input_blocks.11.0.in_layers.2.weight': 'down_blocks.3.resnets.1.conv1.weight',
    'model.diffusion_model.input_blocks.11.0.in_layers.2.bias': 'down_blocks.3.resnets.1.conv1.bias',
    'model.diffusion_model.input_blocks.11.0.emb_layers.1.weight': 'down_blocks.3.resnets.1.time_emb_proj.weight',
    'model.diffusion_model.input_blocks.11.0.emb_layers.1.bias': 'down_blocks.3.resnets.1.time_emb_proj.bias',
    'model.diffusion_model.input_blocks.11.0.out_layers.0.weight': 'down_blocks.3.resnets.1.norm2.weight',
    'model.diffusion_model.input_blocks.11.0.out_layers.0.bias': 'down_blocks.3.resnets.1.norm2.bias',
    'model.diffusion_model.input_blocks.11.0.out_layers.3.weight': 'down_blocks.3.resnets.1.conv2.weight',
    'model.diffusion_model.input_blocks.11.0.out_layers.3.bias': 'down_blocks.3.resnets.1.conv2.bias',
    'model.diffusion_model.middle_block.0.in_layers.0.weight': 'mid_block.resnets.0.norm1.weight',
    'model.diffusion_model.middle_block.0.in_layers.0.bias': 'mid_block.resnets.0.norm1.bias',
    'model.diffusion_model.middle_block.0.in_layers.2.weight': 'mid_block.resnets.0.conv1.weight',
    'model.diffusion_model.middle_block.0.in_layers.2.bias': 'mid_block.resnets.0.conv1.bias',
    'model.diffusion_model.middle_block.0.emb_layers.1.weight': 'mid_block.resnets.0.time_emb_proj.weight',
    'model.diffusion_model.middle_block.0.emb_layers.1.bias': 'mid_block.resnets.0.time_emb_proj.bias',
    'model.diffusion_model.middle_block.0.out_layers.0.weight': 'mid_block.resnets.0.norm2.weight',
    'model.diffusion_model.middle_block.0.out_layers.0.bias': 'mid_block.resnets.0.norm2.bias',
    'model.diffusion_model.middle_block.0.out_layers.3.weight': 'mid_block.resnets.0.conv2.weight',
    'model.diffusion_model.middle_block.0.out_layers.3.bias': 'mid_block.resnets.0.conv2.bias',
    'model.diffusion_model.middle_block.1.norm.weight': 'mid_block.attentions.0.norm.weight',
    'model.diffusion_model.middle_block.1.norm.bias': 'mid_block.attentions.0.norm.bias',
    'model.diffusion_model.middle_block.1.proj_in.weight': 'mid_block.attentions.0.proj_in.weight',
    'model.diffusion_model.middle_block.1.proj_in.bias': 'mid_block.attentions.0.proj_in.bias',
    'model.diffusion_model.middle_block.1.transformer_blocks.0.norm1.weight': 'mid_block.attentions.0.transformer_blocks.0.norm1.weight',
    'model.diffusion_model.middle_block.1.transformer_blocks.0.norm1.bias': 'mid_block.attentions.0.transformer_blocks.0.norm1.bias',
    'model.diffusion_model.middle_block.1.transformer_blocks.0.attn1.to_q.weight': 'mid_block.attentions.0.transformer_blocks.0.attn1.to_q.weight',
    'model.diffusion_model.middle_block.1.transformer_blocks.0.attn1.to_k.weight': 'mid_block.attentions.0.transformer_blocks.0.attn1.to_k.weight',
    'model.diffusion_model.middle_block.1.transformer_blocks.0.attn1.to_v.weight': 'mid_block.attentions.0.transformer_blocks.0.attn1.to_v.weight',
    'model.diffusion_model.middle_block.1.transformer_blocks.0.attn1.to_out.0.weight': 'mid_block.attentions.0.transformer_blocks.0.attn1.to_out.0.weight',
    'model.diffusion_model.middle_block.1.transformer_blocks.0.attn1.to_out.0.bias': 'mid_block.attentions.0.transformer_blocks.0.attn1.to_out.0.bias',
    'model.diffusion_model.middle_block.1.transformer_blocks.0.norm2.weight': 'mid_block.attentions.0.transformer_blocks.0.norm2.weight',
    'model.diffusion_model.middle_block.1.transformer_blocks.0.norm2.bias': 'mid_block.attentions.0.transformer_blocks.0.norm2.bias',
    'model.diffusion_model.middle_block.1.transformer_blocks.0.attn2.to_q.weight': 'mid_block.attentions.0.transformer_blocks.0.attn2.to_q.weight',
    'model.diffusion_model.middle_block.1.transformer_blocks.0.attn2.to_k.weight': 'mid_block.attentions.0.transformer_blocks.0.attn2.to_k.weight',
    'model.diffusion_model.middle_block.1.transformer_blocks.0.attn2.to_v.weight': 'mid_block.attentions.0.transformer_blocks.0.attn2.to_v.weight',
    'model.diffusion_model.middle_block.1.transformer_blocks.0.attn2.to_out.0.weight': 'mid_block.attentions.0.transformer_blocks.0.attn2.to_out.0.weight',
    'model.diffusion_model.middle_block.1.transformer_blocks.0.attn2.to_out.0.bias': 'mid_block.attentions.0.transformer_blocks.0.attn2.to_out.0.bias',
    'model.diffusion_model.middle_block.1.transformer_blocks.0.norm3.weight': 'mid_block.attentions.0.transformer_blocks.0.norm3.weight',
    'model.diffusion_model.middle_block.1.transformer_blocks.0.norm3.bias': 'mid_block.attentions.0.transformer_blocks.0.norm3.bias',
    'model.diffusion_model.middle_block.1.transformer_blocks.0.ff.net.0.proj.weight': 'mid_block.attentions.0.transformer_blocks.0.ff.net.0.proj.weight',
    'model.diffusion_model.middle_block.1.transformer_blocks.0.ff.net.0.proj.bias': 'mid_block.attentions.0.transformer_blocks.0.ff.net.0.proj.bias',
    'model.diffusion_model.middle_block.1.transformer_blocks.0.ff.net.2.weight': 'mid_block.attentions.0.transformer_blocks.0.ff.net.2.weight',
    'model.diffusion_model.middle_block.1.transformer_blocks.0.ff.net.2.bias': 'mid_block.attentions.0.transformer_blocks.0.ff.net.2.bias',
    'model.diffusion_model.middle_block.1.proj_out.weight': 'mid_block.attentions.0.proj_out.weight',
    'model.diffusion_model.middle_block.1.proj_out.bias': 'mid_block.attentions.0.proj_out.bias',
    'model.diffusion_model.middle_block.2.in_layers.0.weight': 'mid_block.resnets.1.norm1.weight',
    'model.diffusion_model.middle_block.2.in_layers.0.bias': 'mid_block.resnets.1.norm1.bias',
    'model.diffusion_model.middle_block.2.in_layers.2.weight': 'mid_block.resnets.1.conv1.weight',
    'model.diffusion_model.middle_block.2.in_layers.2.bias': 'mid_block.resnets.1.conv1.bias',
    'model.diffusion_model.middle_block.2.emb_layers.1.weight': 'mid_block.resnets.1.time_emb_proj.weight',
    'model.diffusion_model.middle_block.2.emb_layers.1.bias': 'mid_block.resnets.1.time_emb_proj.bias',
    'model.diffusion_model.middle_block.2.out_layers.0.weight': 'mid_block.resnets.1.norm2.weight',
    'model.diffusion_model.middle_block.2.out_layers.0.bias': 'mid_block.resnets.1.norm2.bias',
    'model.diffusion_model.middle_block.2.out_layers.3.weight': 'mid_block.resnets.1.conv2.weight',
    'model.diffusion_model.middle_block.2.out_layers.3.bias': 'mid_block.resnets.1.conv2.bias',
    'model.diffusion_model.output_blocks.0.0.in_layers.0.weight': 'up_blocks.0.resnets.0.norm1.weight',
    'model.diffusion_model.output_blocks.0.0.in_layers.0.bias': 'up_blocks.0.resnets.0.norm1.bias',
    'model.diffusion_model.output_blocks.0.0.in_layers.2.weight': 'up_blocks.0.resnets.0.conv1.weight',
    'model.diffusion_model.output_blocks.0.0.in_layers.2.bias': 'up_blocks.0.resnets.0.conv1.bias',
    'model.diffusion_model.output_blocks.0.0.emb_layers.1.weight': 'up_blocks.0.resnets.0.time_emb_proj.weight',
    'model.diffusion_model.output_blocks.0.0.emb_layers.1.bias': 'up_blocks.0.resnets.0.time_emb_proj.bias',
    'model.diffusion_model.output_blocks.0.0.out_layers.0.weight': 'up_blocks.0.resnets.0.norm2.weight',
    'model.diffusion_model.output_blocks.0.0.out_layers.0.bias': 'up_blocks.0.resnets.0.norm2.bias',
    'model.diffusion_model.output_blocks.0.0.out_layers.3.weight': 'up_blocks.0.resnets.0.conv2.weight',
    'model.diffusion_model.output_blocks.0.0.out_layers.3.bias': 'up_blocks.0.resnets.0.conv2.bias',
    'model.diffusion_model.output_blocks.0.0.skip_connection.weight': 'up_blocks.0.resnets.0.conv_shortcut.weight',
    'model.diffusion_model.output_blocks.0.0.skip_connection.bias': 'up_blocks.0.resnets.0.conv_shortcut.bias',
    'model.diffusion_model.output_blocks.1.0.in_layers.0.weight': 'up_blocks.0.resnets.1.norm1.weight',
    'model.diffusion_model.output_blocks.1.0.in_layers.0.bias': 'up_blocks.0.resnets.1.norm1.bias',
    'model.diffusion_model.output_blocks.1.0.in_layers.2.weight': 'up_blocks.0.resnets.1.conv1.weight',
    'model.diffusion_model.output_blocks.1.0.in_layers.2.bias': 'up_blocks.0.resnets.1.conv1.bias',
    'model.diffusion_model.output_blocks.1.0.emb_layers.1.weight': 'up_blocks.0.resnets.1.time_emb_proj.weight',
    'model.diffusion_model.output_blocks.1.0.emb_layers.1.bias': 'up_blocks.0.resnets.1.time_emb_proj.bias',
    'model.diffusion_model.output_blocks.1.0.out_layers.0.weight': 'up_blocks.0.resnets.1.norm2.weight',
    'model.diffusion_model.output_blocks.1.0.out_layers.0.bias': 'up_blocks.0.resnets.1.norm2.bias',
    'model.diffusion_model.output_blocks.1.0.out_layers.3.weight': 'up_blocks.0.resnets.1.conv2.weight',
    'model.diffusion_model.output_blocks.1.0.out_layers.3.bias': 'up_blocks.0.resnets.1.conv2.bias',
    'model.diffusion_model.output_blocks.1.0.skip_connection.weight': 'up_blocks.0.resnets.1.conv_shortcut.weight',
    'model.diffusion_model.output_blocks.1.0.skip_connection.bias': 'up_blocks.0.resnets.1.conv_shortcut.bias',
    'model.diffusion_model.output_blocks.2.0.in_layers.0.weight': 'up_blocks.0.resnets.2.norm1.weight',
    'model.diffusion_model.output_blocks.2.0.in_layers.0.bias': 'up_blocks.0.resnets.2.norm1.bias',
    'model.diffusion_model.output_blocks.2.0.in_layers.2.weight': 'up_blocks.0.resnets.2.conv1.weight',
    'model.diffusion_model.output_blocks.2.0.in_layers.2.bias': 'up_blocks.0.resnets.2.conv1.bias',
    'model.diffusion_model.output_blocks.2.0.emb_layers.1.weight': 'up_blocks.0.resnets.2.time_emb_proj.weight',
    'model.diffusion_model.output_blocks.2.0.emb_layers.1.bias': 'up_blocks.0.resnets.2.time_emb_proj.bias',
    'model.diffusion_model.output_blocks.2.0.out_layers.0.weight': 'up_blocks.0.resnets.2.norm2.weight',
    'model.diffusion_model.output_blocks.2.0.out_layers.0.bias': 'up_blocks.0.resnets.2.norm2.bias',
    'model.diffusion_model.output_blocks.2.0.out_layers.3.weight': 'up_blocks.0.resnets.2.conv2.weight',
    'model.diffusion_model.output_blocks.2.0.out_layers.3.bias': 'up_blocks.0.resnets.2.conv2.bias',
    'model.diffusion_model.output_blocks.2.0.skip_connection.weight': 'up_blocks.0.resnets.2.conv_shortcut.weight',
    'model.diffusion_model.output_blocks.2.0.skip_connection.bias': 'up_blocks.0.resnets.2.conv_shortcut.bias',
    'model.diffusion_model.output_blocks.2.1.conv.weight': 'up_blocks.0.upsamplers.0.conv.weight',
    'model.diffusion_model.output_blocks.2.1.conv.bias': 'up_blocks.0.upsamplers.0.conv.bias',
    'model.diffusion_model.output_blocks.3.0.in_layers.0.weight': 'up_blocks.1.resnets.0.norm1.weight',
    'model.diffusion_model.output_blocks.3.0.in_layers.0.bias': 'up_blocks.1.resnets.0.norm1.bias',
    'model.diffusion_model.output_blocks.3.0.in_layers.2.weight': 'up_blocks.1.resnets.0.conv1.weight',
    'model.diffusion_model.output_blocks.3.0.in_layers.2.bias': 'up_blocks.1.resnets.0.conv1.bias',
    'model.diffusion_model.output_blocks.3.0.emb_layers.1.weight': 'up_blocks.1.resnets.0.time_emb_proj.weight',
    'model.diffusion_model.output_blocks.3.0.emb_layers.1.bias': 'up_blocks.1.resnets.0.time_emb_proj.bias',
    'model.diffusion_model.output_blocks.3.0.out_layers.0.weight': 'up_blocks.1.resnets.0.norm2.weight',
    'model.diffusion_model.output_blocks.3.0.out_layers.0.bias': 'up_blocks.1.resnets.0.norm2.bias',
    'model.diffusion_model.output_blocks.3.0.out_layers.3.weight': 'up_blocks.1.resnets.0.conv2.weight',
    'model.diffusion_model.output_blocks.3.0.out_layers.3.bias': 'up_blocks.1.resnets.0.conv2.bias',
    'model.diffusion_model.output_blocks.3.0.skip_connection.weight': 'up_blocks.1.resnets.0.conv_shortcut.weight',
    'model.diffusion_model.output_blocks.3.0.skip_connection.bias': 'up_blocks.1.resnets.0.conv_shortcut.bias',
    'model.diffusion_model.output_blocks.3.1.norm.weight': 'up_blocks.1.attentions.0.norm.weight',
    'model.diffusion_model.output_blocks.3.1.norm.bias': 'up_blocks.1.attentions.0.norm.bias',
    'model.diffusion_model.output_blocks.3.1.proj_in.weight': 'up_blocks.1.attentions.0.proj_in.weight',
    'model.diffusion_model.output_blocks.3.1.proj_in.bias': 'up_blocks.1.attentions.0.proj_in.bias',
    'model.diffusion_model.output_blocks.3.1.transformer_blocks.0.norm1.weight': 'up_blocks.1.attentions.0.transformer_blocks.0.norm1.weight',
    'model.diffusion_model.output_blocks.3.1.transformer_blocks.0.norm1.bias': 'up_blocks.1.attentions.0.transformer_blocks.0.norm1.bias',
    'model.diffusion_model.output_blocks.3.1.transformer_blocks.0.attn1.to_q.weight': 'up_blocks.1.attentions.0.transformer_blocks.0.attn1.to_q.weight',
    'model.diffusion_model.output_blocks.3.1.transformer_blocks.0.attn1.to_k.weight': 'up_blocks.1.attentions.0.transformer_blocks.0.attn1.to_k.weight',
    'model.diffusion_model.output_blocks.3.1.transformer_blocks.0.attn1.to_v.weight': 'up_blocks.1.attentions.0.transformer_blocks.0.attn1.to_v.weight',
    'model.diffusion_model.output_blocks.3.1.transformer_blocks.0.attn1.to_out.0.weight': 'up_blocks.1.attentions.0.transformer_blocks.0.attn1.to_out.0.weight',
    'model.diffusion_model.output_blocks.3.1.transformer_blocks.0.attn1.to_out.0.bias': 'up_blocks.1.attentions.0.transformer_blocks.0.attn1.to_out.0.bias',
    'model.diffusion_model.output_blocks.3.1.transformer_blocks.0.norm2.weight': 'up_blocks.1.attentions.0.transformer_blocks.0.norm2.weight',
    'model.diffusion_model.output_blocks.3.1.transformer_blocks.0.norm2.bias': 'up_blocks.1.attentions.0.transformer_blocks.0.norm2.bias',
    'model.diffusion_model.output_blocks.3.1.transformer_blocks.0.attn2.to_q.weight': 'up_blocks.1.attentions.0.transformer_blocks.0.attn2.to_q.weight',
    'model.diffusion_model.output_blocks.3.1.transformer_blocks.0.attn2.to_k.weight': 'up_blocks.1.attentions.0.transformer_blocks.0.attn2.to_k.weight',
    'model.diffusion_model.output_blocks.3.1.transformer_blocks.0.attn2.to_v.weight': 'up_blocks.1.attentions.0.transformer_blocks.0.attn2.to_v.weight',
    'model.diffusion_model.output_blocks.3.1.transformer_blocks.0.attn2.to_out.0.weight': 'up_blocks.1.attentions.0.transformer_blocks.0.attn2.to_out.0.weight',
    'model.diffusion_model.output_blocks.3.1.transformer_blocks.0.attn2.to_out.0.bias': 'up_blocks.1.attentions.0.transformer_blocks.0.attn2.to_out.0.bias',
    'model.diffusion_model.output_blocks.3.1.transformer_blocks.0.norm3.weight': 'up_blocks.1.attentions.0.transformer_blocks.0.norm3.weight',
    'model.diffusion_model.output_blocks.3.1.transformer_blocks.0.norm3.bias': 'up_blocks.1.attentions.0.transformer_blocks.0.norm3.bias',
    'model.diffusion_model.output_blocks.3.1.transformer_blocks.0.ff.net.0.proj.weight': 'up_blocks.1.attentions.0.transformer_blocks.0.ff.net.0.proj.weight',
    'model.diffusion_model.output_blocks.3.1.transformer_blocks.0.ff.net.0.proj.bias': 'up_blocks.1.attentions.0.transformer_blocks.0.ff.net.0.proj.bias',
    'model.diffusion_model.output_blocks.3.1.transformer_blocks.0.ff.net.2.weight': 'up_blocks.1.attentions.0.transformer_blocks.0.ff.net.2.weight',
    'model.diffusion_model.output_blocks.3.1.transformer_blocks.0.ff.net.2.bias': 'up_blocks.1.attentions.0.transformer_blocks.0.ff.net.2.bias',
    'model.diffusion_model.output_blocks.3.1.proj_out.weight': 'up_blocks.1.attentions.0.proj_out.weight',
    'model.diffusion_model.output_blocks.3.1.proj_out.bias': 'up_blocks.1.attentions.0.proj_out.bias',
    'model.diffusion_model.output_blocks.4.0.in_layers.0.weight': 'up_blocks.1.resnets.1.norm1.weight',
    'model.diffusion_model.output_blocks.4.0.in_layers.0.bias': 'up_blocks.1.resnets.1.norm1.bias',
    'model.diffusion_model.output_blocks.4.0.in_layers.2.weight': 'up_blocks.1.resnets.1.conv1.weight',
    'model.diffusion_model.output_blocks.4.0.in_layers.2.bias': 'up_blocks.1.resnets.1.conv1.bias',
    'model.diffusion_model.output_blocks.4.0.emb_layers.1.weight': 'up_blocks.1.resnets.1.time_emb_proj.weight',
    'model.diffusion_model.output_blocks.4.0.emb_layers.1.bias': 'up_blocks.1.resnets.1.time_emb_proj.bias',
    'model.diffusion_model.output_blocks.4.0.out_layers.0.weight': 'up_blocks.1.resnets.1.norm2.weight',
    'model.diffusion_model.output_blocks.4.0.out_layers.0.bias': 'up_blocks.1.resnets.1.norm2.bias',
    'model.diffusion_model.output_blocks.4.0.out_layers.3.weight': 'up_blocks.1.resnets.1.conv2.weight',
    'model.diffusion_model.output_blocks.4.0.out_layers.3.bias': 'up_blocks.1.resnets.1.conv2.bias',
    'model.diffusion_model.output_blocks.4.0.skip_connection.weight': 'up_blocks.1.resnets.1.conv_shortcut.weight',
    'model.diffusion_model.output_blocks.4.0.skip_connection.bias': 'up_blocks.1.resnets.1.conv_shortcut.bias',
    'model.diffusion_model.output_blocks.4.1.norm.weight': 'up_blocks.1.attentions.1.norm.weight',
    'model.diffusion_model.output_blocks.4.1.norm.bias': 'up_blocks.1.attentions.1.norm.bias',
    'model.diffusion_model.output_blocks.4.1.proj_in.weight': 'up_blocks.1.attentions.1.proj_in.weight',
    'model.diffusion_model.output_blocks.4.1.proj_in.bias': 'up_blocks.1.attentions.1.proj_in.bias',
    'model.diffusion_model.output_blocks.4.1.transformer_blocks.0.norm1.weight': 'up_blocks.1.attentions.1.transformer_blocks.0.norm1.weight',
    'model.diffusion_model.output_blocks.4.1.transformer_blocks.0.norm1.bias': 'up_blocks.1.attentions.1.transformer_blocks.0.norm1.bias',
    'model.diffusion_model.output_blocks.4.1.transformer_blocks.0.attn1.to_q.weight': 'up_blocks.1.attentions.1.transformer_blocks.0.attn1.to_q.weight',
    'model.diffusion_model.output_blocks.4.1.transformer_blocks.0.attn1.to_k.weight': 'up_blocks.1.attentions.1.transformer_blocks.0.attn1.to_k.weight',
    'model.diffusion_model.output_blocks.4.1.transformer_blocks.0.attn1.to_v.weight': 'up_blocks.1.attentions.1.transformer_blocks.0.attn1.to_v.weight',
    'model.diffusion_model.output_blocks.4.1.transformer_blocks.0.attn1.to_out.0.weight': 'up_blocks.1.attentions.1.transformer_blocks.0.attn1.to_out.0.weight',
    'model.diffusion_model.output_blocks.4.1.transformer_blocks.0.attn1.to_out.0.bias': 'up_blocks.1.attentions.1.transformer_blocks.0.attn1.to_out.0.bias',
    'model.diffusion_model.output_blocks.4.1.transformer_blocks.0.norm2.weight': 'up_blocks.1.attentions.1.transformer_blocks.0.norm2.weight',
    'model.diffusion_model.output_blocks.4.1.transformer_blocks.0.norm2.bias': 'up_blocks.1.attentions.1.transformer_blocks.0.norm2.bias',
    'model.diffusion_model.output_blocks.4.1.transformer_blocks.0.attn2.to_q.weight': 'up_blocks.1.attentions.1.transformer_blocks.0.attn2.to_q.weight',
    'model.diffusion_model.output_blocks.4.1.transformer_blocks.0.attn2.to_k.weight': 'up_blocks.1.attentions.1.transformer_blocks.0.attn2.to_k.weight',
    'model.diffusion_model.output_blocks.4.1.transformer_blocks.0.attn2.to_v.weight': 'up_blocks.1.attentions.1.transformer_blocks.0.attn2.to_v.weight',
    'model.diffusion_model.output_blocks.4.1.transformer_blocks.0.attn2.to_out.0.weight': 'up_blocks.1.attentions.1.transformer_blocks.0.attn2.to_out.0.weight',
    'model.diffusion_model.output_blocks.4.1.transformer_blocks.0.attn2.to_out.0.bias': 'up_blocks.1.attentions.1.transformer_blocks.0.attn2.to_out.0.bias',
    'model.diffusion_model.output_blocks.4.1.transformer_blocks.0.norm3.weight': 'up_blocks.1.attentions.1.transformer_blocks.0.norm3.weight',
    'model.diffusion_model.output_blocks.4.1.transformer_blocks.0.norm3.bias': 'up_blocks.1.attentions.1.transformer_blocks.0.norm3.bias',
    'model.diffusion_model.output_blocks.4.1.transformer_blocks.0.ff.net.0.proj.weight': 'up_blocks.1.attentions.1.transformer_blocks.0.ff.net.0.proj.weight',
    'model.diffusion_model.output_blocks.4.1.transformer_blocks.0.ff.net.0.proj.bias': 'up_blocks.1.attentions.1.transformer_blocks.0.ff.net.0.proj.bias',
    'model.diffusion_model.output_blocks.4.1.transformer_blocks.0.ff.net.2.weight': 'up_blocks.1.attentions.1.transformer_blocks.0.ff.net.2.weight',
    'model.diffusion_model.output_blocks.4.1.transformer_blocks.0.ff.net.2.bias': 'up_blocks.1.attentions.1.transformer_blocks.0.ff.net.2.bias',
    'model.diffusion_model.output_blocks.4.1.proj_out.weight': 'up_blocks.1.attentions.1.proj_out.weight',
    'model.diffusion_model.output_blocks.4.1.proj_out.bias': 'up_blocks.1.attentions.1.proj_out.bias',
    'model.diffusion_model.output_blocks.5.0.in_layers.0.weight': 'up_blocks.1.resnets.2.norm1.weight',
    'model.diffusion_model.output_blocks.5.0.in_layers.0.bias': 'up_blocks.1.resnets.2.norm1.bias',
    'model.diffusion_model.output_blocks.5.0.in_layers.2.weight': 'up_blocks.1.resnets.2.conv1.weight',
    'model.diffusion_model.output_blocks.5.0.in_layers.2.bias': 'up_blocks.1.resnets.2.conv1.bias',
    'model.diffusion_model.output_blocks.5.0.emb_layers.1.weight': 'up_blocks.1.resnets.2.time_emb_proj.weight',
    'model.diffusion_model.output_blocks.5.0.emb_layers.1.bias': 'up_blocks.1.resnets.2.time_emb_proj.bias',
    'model.diffusion_model.output_blocks.5.0.out_layers.0.weight': 'up_blocks.1.resnets.2.norm2.weight',
    'model.diffusion_model.output_blocks.5.0.out_layers.0.bias': 'up_blocks.1.resnets.2.norm2.bias',
    'model.diffusion_model.output_blocks.5.0.out_layers.3.weight': 'up_blocks.1.resnets.2.conv2.weight',
    'model.diffusion_model.output_blocks.5.0.out_layers.3.bias': 'up_blocks.1.resnets.2.conv2.bias',
    'model.diffusion_model.output_blocks.5.0.skip_connection.weight': 'up_blocks.1.resnets.2.conv_shortcut.weight',
    'model.diffusion_model.output_blocks.5.0.skip_connection.bias': 'up_blocks.1.resnets.2.conv_shortcut.bias',
    'model.diffusion_model.output_blocks.5.1.norm.weight': 'up_blocks.1.attentions.2.norm.weight',
    'model.diffusion_model.output_blocks.5.1.norm.bias': 'up_blocks.1.attentions.2.norm.bias',
    'model.diffusion_model.output_blocks.5.1.proj_in.weight': 'up_blocks.1.attentions.2.proj_in.weight',
    'model.diffusion_model.output_blocks.5.1.proj_in.bias': 'up_blocks.1.attentions.2.proj_in.bias',
    'model.diffusion_model.output_blocks.5.1.transformer_blocks.0.norm1.weight': 'up_blocks.1.attentions.2.transformer_blocks.0.norm1.weight',
    'model.diffusion_model.output_blocks.5.1.transformer_blocks.0.norm1.bias': 'up_blocks.1.attentions.2.transformer_blocks.0.norm1.bias',
    'model.diffusion_model.output_blocks.5.1.transformer_blocks.0.attn1.to_q.weight': 'up_blocks.1.attentions.2.transformer_blocks.0.attn1.to_q.weight',
    'model.diffusion_model.output_blocks.5.1.transformer_blocks.0.attn1.to_k.weight': 'up_blocks.1.attentions.2.transformer_blocks.0.attn1.to_k.weight',
    'model.diffusion_model.output_blocks.5.1.transformer_blocks.0.attn1.to_v.weight': 'up_blocks.1.attentions.2.transformer_blocks.0.attn1.to_v.weight',
    'model.diffusion_model.output_blocks.5.1.transformer_blocks.0.attn1.to_out.0.weight': 'up_blocks.1.attentions.2.transformer_blocks.0.attn1.to_out.0.weight',
    'model.diffusion_model.output_blocks.5.1.transformer_blocks.0.attn1.to_out.0.bias': 'up_blocks.1.attentions.2.transformer_blocks.0.attn1.to_out.0.bias',
    'model.diffusion_model.output_blocks.5.1.transformer_blocks.0.norm2.weight': 'up_blocks.1.attentions.2.transformer_blocks.0.norm2.weight',
    'model.diffusion_model.output_blocks.5.1.transformer_blocks.0.norm2.bias': 'up_blocks.1.attentions.2.transformer_blocks.0.norm2.bias',
    'model.diffusion_model.output_blocks.5.1.transformer_blocks.0.attn2.to_q.weight': 'up_blocks.1.attentions.2.transformer_blocks.0.attn2.to_q.weight',
    'model.diffusion_model.output_blocks.5.1.transformer_blocks.0.attn2.to_k.weight': 'up_blocks.1.attentions.2.transformer_blocks.0.attn2.to_k.weight',
    'model.diffusion_model.output_blocks.5.1.transformer_blocks.0.attn2.to_v.weight': 'up_blocks.1.attentions.2.transformer_blocks.0.attn2.to_v.weight',
    'model.diffusion_model.output_blocks.5.1.transformer_blocks.0.attn2.to_out.0.weight': 'up_blocks.1.attentions.2.transformer_blocks.0.attn2.to_out.0.weight',
    'model.diffusion_model.output_blocks.5.1.transformer_blocks.0.attn2.to_out.0.bias': 'up_blocks.1.attentions.2.transformer_blocks.0.attn2.to_out.0.bias',
    'model.diffusion_model.output_blocks.5.1.transformer_blocks.0.norm3.weight': 'up_blocks.1.attentions.2.transformer_blocks.0.norm3.weight',
    'model.diffusion_model.output_blocks.5.1.transformer_blocks.0.norm3.bias': 'up_blocks.1.attentions.2.transformer_blocks.0.norm3.bias',
    'model.diffusion_model.output_blocks.5.1.transformer_blocks.0.ff.net.0.proj.weight': 'up_blocks.1.attentions.2.transformer_blocks.0.ff.net.0.proj.weight',
    'model.diffusion_model.output_blocks.5.1.transformer_blocks.0.ff.net.0.proj.bias': 'up_blocks.1.attentions.2.transformer_blocks.0.ff.net.0.proj.bias',
    'model.diffusion_model.output_blocks.5.1.transformer_blocks.0.ff.net.2.weight': 'up_blocks.1.attentions.2.transformer_blocks.0.ff.net.2.weight',
    'model.diffusion_model.output_blocks.5.1.transformer_blocks.0.ff.net.2.bias': 'up_blocks.1.attentions.2.transformer_blocks.0.ff.net.2.bias',
    'model.diffusion_model.output_blocks.5.1.proj_out.weight': 'up_blocks.1.attentions.2.proj_out.weight',
    'model.diffusion_model.output_blocks.5.1.proj_out.bias': 'up_blocks.1.attentions.2.proj_out.bias',
    'model.diffusion_model.output_blocks.5.2.conv.weight': 'up_blocks.1.upsamplers.0.conv.weight',
    'model.diffusion_model.output_blocks.5.2.conv.bias': 'up_blocks.1.upsamplers.0.conv.bias',
    'model.diffusion_model.output_blocks.6.0.in_layers.0.weight': 'up_blocks.2.resnets.0.norm1.weight',
    'model.diffusion_model.output_blocks.6.0.in_layers.0.bias': 'up_blocks.2.resnets.0.norm1.bias',
    'model.diffusion_model.output_blocks.6.0.in_layers.2.weight': 'up_blocks.2.resnets.0.conv1.weight',
    'model.diffusion_model.output_blocks.6.0.in_layers.2.bias': 'up_blocks.2.resnets.0.conv1.bias',
    'model.diffusion_model.output_blocks.6.0.emb_layers.1.weight': 'up_blocks.2.resnets.0.time_emb_proj.weight',
    'model.diffusion_model.output_blocks.6.0.emb_layers.1.bias': 'up_blocks.2.resnets.0.time_emb_proj.bias',
    'model.diffusion_model.output_blocks.6.0.out_layers.0.weight': 'up_blocks.2.resnets.0.norm2.weight',
    'model.diffusion_model.output_blocks.6.0.out_layers.0.bias': 'up_blocks.2.resnets.0.norm2.bias',
    'model.diffusion_model.output_blocks.6.0.out_layers.3.weight': 'up_blocks.2.resnets.0.conv2.weight',
    'model.diffusion_model.output_blocks.6.0.out_layers.3.bias': 'up_blocks.2.resnets.0.conv2.bias',
    'model.diffusion_model.output_blocks.6.0.skip_connection.weight': 'up_blocks.2.resnets.0.conv_shortcut.weight',
    'model.diffusion_model.output_blocks.6.0.skip_connection.bias': 'up_blocks.2.resnets.0.conv_shortcut.bias',
    'model.diffusion_model.output_blocks.6.1.norm.weight': 'up_blocks.2.attentions.0.norm.weight',
    'model.diffusion_model.output_blocks.6.1.norm.bias': 'up_blocks.2.attentions.0.norm.bias',
    'model.diffusion_model.output_blocks.6.1.proj_in.weight': 'up_blocks.2.attentions.0.proj_in.weight',
    'model.diffusion_model.output_blocks.6.1.proj_in.bias': 'up_blocks.2.attentions.0.proj_in.bias',
    'model.diffusion_model.output_blocks.6.1.transformer_blocks.0.norm1.weight': 'up_blocks.2.attentions.0.transformer_blocks.0.norm1.weight',
    'model.diffusion_model.output_blocks.6.1.transformer_blocks.0.norm1.bias': 'up_blocks.2.attentions.0.transformer_blocks.0.norm1.bias',
    'model.diffusion_model.output_blocks.6.1.transformer_blocks.0.attn1.to_q.weight': 'up_blocks.2.attentions.0.transformer_blocks.0.attn1.to_q.weight',
    'model.diffusion_model.output_blocks.6.1.transformer_blocks.0.attn1.to_k.weight': 'up_blocks.2.attentions.0.transformer_blocks.0.attn1.to_k.weight',
    'model.diffusion_model.output_blocks.6.1.transformer_blocks.0.attn1.to_v.weight': 'up_blocks.2.attentions.0.transformer_blocks.0.attn1.to_v.weight',
    'model.diffusion_model.output_blocks.6.1.transformer_blocks.0.attn1.to_out.0.weight': 'up_blocks.2.attentions.0.transformer_blocks.0.attn1.to_out.0.weight',
    'model.diffusion_model.output_blocks.6.1.transformer_blocks.0.attn1.to_out.0.bias': 'up_blocks.2.attentions.0.transformer_blocks.0.attn1.to_out.0.bias',
    'model.diffusion_model.output_blocks.6.1.transformer_blocks.0.norm2.weight': 'up_blocks.2.attentions.0.transformer_blocks.0.norm2.weight',
    'model.diffusion_model.output_blocks.6.1.transformer_blocks.0.norm2.bias': 'up_blocks.2.attentions.0.transformer_blocks.0.norm2.bias',
    'model.diffusion_model.output_blocks.6.1.transformer_blocks.0.attn2.to_q.weight': 'up_blocks.2.attentions.0.transformer_blocks.0.attn2.to_q.weight',
    'model.diffusion_model.output_blocks.6.1.transformer_blocks.0.attn2.to_k.weight': 'up_blocks.2.attentions.0.transformer_blocks.0.attn2.to_k.weight',
    'model.diffusion_model.output_blocks.6.1.transformer_blocks.0.attn2.to_v.weight': 'up_blocks.2.attentions.0.transformer_blocks.0.attn2.to_v.weight',
    'model.diffusion_model.output_blocks.6.1.transformer_blocks.0.attn2.to_out.0.weight': 'up_blocks.2.attentions.0.transformer_blocks.0.attn2.to_out.0.weight',
    'model.diffusion_model.output_blocks.6.1.transformer_blocks.0.attn2.to_out.0.bias': 'up_blocks.2.attentions.0.transformer_blocks.0.attn2.to_out.0.bias',
    'model.diffusion_model.output_blocks.6.1.transformer_blocks.0.norm3.weight': 'up_blocks.2.attentions.0.transformer_blocks.0.norm3.weight',
    'model.diffusion_model.output_blocks.6.1.transformer_blocks.0.norm3.bias': 'up_blocks.2.attentions.0.transformer_blocks.0.norm3.bias',
    'model.diffusion_model.output_blocks.6.1.transformer_blocks.0.ff.net.0.proj.weight': 'up_blocks.2.attentions.0.transformer_blocks.0.ff.net.0.proj.weight',
    'model.diffusion_model.output_blocks.6.1.transformer_blocks.0.ff.net.0.proj.bias': 'up_blocks.2.attentions.0.transformer_blocks.0.ff.net.0.proj.bias',
    'model.diffusion_model.output_blocks.6.1.transformer_blocks.0.ff.net.2.weight': 'up_blocks.2.attentions.0.transformer_blocks.0.ff.net.2.weight',
    'model.diffusion_model.output_blocks.6.1.transformer_blocks.0.ff.net.2.bias': 'up_blocks.2.attentions.0.transformer_blocks.0.ff.net.2.bias',
    'model.diffusion_model.output_blocks.6.1.proj_out.weight': 'up_blocks.2.attentions.0.proj_out.weight',
    'model.diffusion_model.output_blocks.6.1.proj_out.bias': 'up_blocks.2.attentions.0.proj_out.bias',
    'model.diffusion_model.output_blocks.7.0.in_layers.0.weight': 'up_blocks.2.resnets.1.norm1.weight',
    'model.diffusion_model.output_blocks.7.0.in_layers.0.bias': 'up_blocks.2.resnets.1.norm1.bias',
    'model.diffusion_model.output_blocks.7.0.in_layers.2.weight': 'up_blocks.2.resnets.1.conv1.weight',
    'model.diffusion_model.output_blocks.7.0.in_layers.2.bias': 'up_blocks.2.resnets.1.conv1.bias',
    'model.diffusion_model.output_blocks.7.0.emb_layers.1.weight': 'up_blocks.2.resnets.1.time_emb_proj.weight',
    'model.diffusion_model.output_blocks.7.0.emb_layers.1.bias': 'up_blocks.2.resnets.1.time_emb_proj.bias',
    'model.diffusion_model.output_blocks.7.0.out_layers.0.weight': 'up_blocks.2.resnets.1.norm2.weight',
    'model.diffusion_model.output_blocks.7.0.out_layers.0.bias': 'up_blocks.2.resnets.1.norm2.bias',
    'model.diffusion_model.output_blocks.7.0.out_layers.3.weight': 'up_blocks.2.resnets.1.conv2.weight',
    'model.diffusion_model.output_blocks.7.0.out_layers.3.bias': 'up_blocks.2.resnets.1.conv2.bias',
    'model.diffusion_model.output_blocks.7.0.skip_connection.weight': 'up_blocks.2.resnets.1.conv_shortcut.weight',
    'model.diffusion_model.output_blocks.7.0.skip_connection.bias': 'up_blocks.2.resnets.1.conv_shortcut.bias',
    'model.diffusion_model.output_blocks.7.1.norm.weight': 'up_blocks.2.attentions.1.norm.weight',
    'model.diffusion_model.output_blocks.7.1.norm.bias': 'up_blocks.2.attentions.1.norm.bias',
    'model.diffusion_model.output_blocks.7.1.proj_in.weight': 'up_blocks.2.attentions.1.proj_in.weight',
    'model.diffusion_model.output_blocks.7.1.proj_in.bias': 'up_blocks.2.attentions.1.proj_in.bias',
    'model.diffusion_model.output_blocks.7.1.transformer_blocks.0.norm1.weight': 'up_blocks.2.attentions.1.transformer_blocks.0.norm1.weight',
    'model.diffusion_model.output_blocks.7.1.transformer_blocks.0.norm1.bias': 'up_blocks.2.attentions.1.transformer_blocks.0.norm1.bias',
    'model.diffusion_model.output_blocks.7.1.transformer_blocks.0.attn1.to_q.weight': 'up_blocks.2.attentions.1.transformer_blocks.0.attn1.to_q.weight',
    'model.diffusion_model.output_blocks.7.1.transformer_blocks.0.attn1.to_k.weight': 'up_blocks.2.attentions.1.transformer_blocks.0.attn1.to_k.weight',
    'model.diffusion_model.output_blocks.7.1.transformer_blocks.0.attn1.to_v.weight': 'up_blocks.2.attentions.1.transformer_blocks.0.attn1.to_v.weight',
    'model.diffusion_model.output_blocks.7.1.transformer_blocks.0.attn1.to_out.0.weight': 'up_blocks.2.attentions.1.transformer_blocks.0.attn1.to_out.0.weight',
    'model.diffusion_model.output_blocks.7.1.transformer_blocks.0.attn1.to_out.0.bias': 'up_blocks.2.attentions.1.transformer_blocks.0.attn1.to_out.0.bias',
    'model.diffusion_model.output_blocks.7.1.transformer_blocks.0.norm2.weight': 'up_blocks.2.attentions.1.transformer_blocks.0.norm2.weight',
    'model.diffusion_model.output_blocks.7.1.transformer_blocks.0.norm2.bias': 'up_blocks.2.attentions.1.transformer_blocks.0.norm2.bias',
    'model.diffusion_model.output_blocks.7.1.transformer_blocks.0.attn2.to_q.weight': 'up_blocks.2.attentions.1.transformer_blocks.0.attn2.to_q.weight',
    'model.diffusion_model.output_blocks.7.1.transformer_blocks.0.attn2.to_k.weight': 'up_blocks.2.attentions.1.transformer_blocks.0.attn2.to_k.weight',
    'model.diffusion_model.output_blocks.7.1.transformer_blocks.0.attn2.to_v.weight': 'up_blocks.2.attentions.1.transformer_blocks.0.attn2.to_v.weight',
    'model.diffusion_model.output_blocks.7.1.transformer_blocks.0.attn2.to_out.0.weight': 'up_blocks.2.attentions.1.transformer_blocks.0.attn2.to_out.0.weight',
    'model.diffusion_model.output_blocks.7.1.transformer_blocks.0.attn2.to_out.0.bias': 'up_blocks.2.attentions.1.transformer_blocks.0.attn2.to_out.0.bias',
    'model.diffusion_model.output_blocks.7.1.transformer_blocks.0.norm3.weight': 'up_blocks.2.attentions.1.transformer_blocks.0.norm3.weight',
    'model.diffusion_model.output_blocks.7.1.transformer_blocks.0.norm3.bias': 'up_blocks.2.attentions.1.transformer_blocks.0.norm3.bias',
    'model.diffusion_model.output_blocks.7.1.transformer_blocks.0.ff.net.0.proj.weight': 'up_blocks.2.attentions.1.transformer_blocks.0.ff.net.0.proj.weight',
    'model.diffusion_model.output_blocks.7.1.transformer_blocks.0.ff.net.0.proj.bias': 'up_blocks.2.attentions.1.transformer_blocks.0.ff.net.0.proj.bias',
    'model.diffusion_model.output_blocks.7.1.transformer_blocks.0.ff.net.2.weight': 'up_blocks.2.attentions.1.transformer_blocks.0.ff.net.2.weight',
    'model.diffusion_model.output_blocks.7.1.transformer_blocks.0.ff.net.2.bias': 'up_blocks.2.attentions.1.transformer_blocks.0.ff.net.2.bias',
    'model.diffusion_model.output_blocks.7.1.proj_out.weight': 'up_blocks.2.attentions.1.proj_out.weight',
    'model.diffusion_model.output_blocks.7.1.proj_out.bias': 'up_blocks.2.attentions.1.proj_out.bias',
    'model.diffusion_model.output_blocks.8.0.in_layers.0.weight': 'up_blocks.2.resnets.2.norm1.weight',
    'model.diffusion_model.output_blocks.8.0.in_layers.0.bias': 'up_blocks.2.resnets.2.norm1.bias',
    'model.diffusion_model.output_blocks.8.0.in_layers.2.weight': 'up_blocks.2.resnets.2.conv1.weight',
    'model.diffusion_model.output_blocks.8.0.in_layers.2.bias': 'up_blocks.2.resnets.2.conv1.bias',
    'model.diffusion_model.output_blocks.8.0.emb_layers.1.weight': 'up_blocks.2.resnets.2.time_emb_proj.weight',
    'model.diffusion_model.output_blocks.8.0.emb_layers.1.bias': 'up_blocks.2.resnets.2.time_emb_proj.bias',
    'model.diffusion_model.output_blocks.8.0.out_layers.0.weight': 'up_blocks.2.resnets.2.norm2.weight',
    'model.diffusion_model.output_blocks.8.0.out_layers.0.bias': 'up_blocks.2.resnets.2.norm2.bias',
    'model.diffusion_model.output_blocks.8.0.out_layers.3.weight': 'up_blocks.2.resnets.2.conv2.weight',
    'model.diffusion_model.output_blocks.8.0.out_layers.3.bias': 'up_blocks.2.resnets.2.conv2.bias',
    'model.diffusion_model.output_blocks.8.0.skip_connection.weight': 'up_blocks.2.resnets.2.conv_shortcut.weight',
    'model.diffusion_model.output_blocks.8.0.skip_connection.bias': 'up_blocks.2.resnets.2.conv_shortcut.bias',
    'model.diffusion_model.output_blocks.8.1.norm.weight': 'up_blocks.2.attentions.2.norm.weight',
    'model.diffusion_model.output_blocks.8.1.norm.bias': 'up_blocks.2.attentions.2.norm.bias',
    'model.diffusion_model.output_blocks.8.1.proj_in.weight': 'up_blocks.2.attentions.2.proj_in.weight',
    'model.diffusion_model.output_blocks.8.1.proj_in.bias': 'up_blocks.2.attentions.2.proj_in.bias',
    'model.diffusion_model.output_blocks.8.1.transformer_blocks.0.norm1.weight': 'up_blocks.2.attentions.2.transformer_blocks.0.norm1.weight',
    'model.diffusion_model.output_blocks.8.1.transformer_blocks.0.norm1.bias': 'up_blocks.2.attentions.2.transformer_blocks.0.norm1.bias',
    'model.diffusion_model.output_blocks.8.1.transformer_blocks.0.attn1.to_q.weight': 'up_blocks.2.attentions.2.transformer_blocks.0.attn1.to_q.weight',
    'model.diffusion_model.output_blocks.8.1.transformer_blocks.0.attn1.to_k.weight': 'up_blocks.2.attentions.2.transformer_blocks.0.attn1.to_k.weight',
    'model.diffusion_model.output_blocks.8.1.transformer_blocks.0.attn1.to_v.weight': 'up_blocks.2.attentions.2.transformer_blocks.0.attn1.to_v.weight',
    'model.diffusion_model.output_blocks.8.1.transformer_blocks.0.attn1.to_out.0.weight': 'up_blocks.2.attentions.2.transformer_blocks.0.attn1.to_out.0.weight',
    'model.diffusion_model.output_blocks.8.1.transformer_blocks.0.attn1.to_out.0.bias': 'up_blocks.2.attentions.2.transformer_blocks.0.attn1.to_out.0.bias',
    'model.diffusion_model.output_blocks.8.1.transformer_blocks.0.norm2.weight': 'up_blocks.2.attentions.2.transformer_blocks.0.norm2.weight',
    'model.diffusion_model.output_blocks.8.1.transformer_blocks.0.norm2.bias': 'up_blocks.2.attentions.2.transformer_blocks.0.norm2.bias',
    'model.diffusion_model.output_blocks.8.1.transformer_blocks.0.attn2.to_q.weight': 'up_blocks.2.attentions.2.transformer_blocks.0.attn2.to_q.weight',
    'model.diffusion_model.output_blocks.8.1.transformer_blocks.0.attn2.to_k.weight': 'up_blocks.2.attentions.2.transformer_blocks.0.attn2.to_k.weight',
    'model.diffusion_model.output_blocks.8.1.transformer_blocks.0.attn2.to_v.weight': 'up_blocks.2.attentions.2.transformer_blocks.0.attn2.to_v.weight',
    'model.diffusion_model.output_blocks.8.1.transformer_blocks.0.attn2.to_out.0.weight': 'up_blocks.2.attentions.2.transformer_blocks.0.attn2.to_out.0.weight',
    'model.diffusion_model.output_blocks.8.1.transformer_blocks.0.attn2.to_out.0.bias': 'up_blocks.2.attentions.2.transformer_blocks.0.attn2.to_out.0.bias',
    'model.diffusion_model.output_blocks.8.1.transformer_blocks.0.norm3.weight': 'up_blocks.2.attentions.2.transformer_blocks.0.norm3.weight',
    'model.diffusion_model.output_blocks.8.1.transformer_blocks.0.norm3.bias': 'up_blocks.2.attentions.2.transformer_blocks.0.norm3.bias',
    'model.diffusion_model.output_blocks.8.1.transformer_blocks.0.ff.net.0.proj.weight': 'up_blocks.2.attentions.2.transformer_blocks.0.ff.net.0.proj.weight',
    'model.diffusion_model.output_blocks.8.1.transformer_blocks.0.ff.net.0.proj.bias': 'up_blocks.2.attentions.2.transformer_blocks.0.ff.net.0.proj.bias',
    'model.diffusion_model.output_blocks.8.1.transformer_blocks.0.ff.net.2.weight': 'up_blocks.2.attentions.2.transformer_blocks.0.ff.net.2.weight',
    'model.diffusion_model.output_blocks.8.1.transformer_blocks.0.ff.net.2.bias': 'up_blocks.2.attentions.2.transformer_blocks.0.ff.net.2.bias',
    'model.diffusion_model.output_blocks.8.1.proj_out.weight': 'up_blocks.2.attentions.2.proj_out.weight',
    'model.diffusion_model.output_blocks.8.1.proj_out.bias': 'up_blocks.2.attentions.2.proj_out.bias',
    'model.diffusion_model.output_blocks.8.2.conv.weight': 'up_blocks.2.upsamplers.0.conv.weight',
    'model.diffusion_model.output_blocks.8.2.conv.bias': 'up_blocks.2.upsamplers.0.conv.bias',
    'model.diffusion_model.output_blocks.9.0.in_layers.0.weight': 'up_blocks.3.resnets.0.norm1.weight',
    'model.diffusion_model.output_blocks.9.0.in_layers.0.bias': 'up_blocks.3.resnets.0.norm1.bias',
    'model.diffusion_model.output_blocks.9.0.in_layers.2.weight': 'up_blocks.3.resnets.0.conv1.weight',
    'model.diffusion_model.output_blocks.9.0.in_layers.2.bias': 'up_blocks.3.resnets.0.conv1.bias',
    'model.diffusion_model.output_blocks.9.0.emb_layers.1.weight': 'up_blocks.3.resnets.0.time_emb_proj.weight',
    'model.diffusion_model.output_blocks.9.0.emb_layers.1.bias': 'up_blocks.3.resnets.0.time_emb_proj.bias',
    'model.diffusion_model.output_blocks.9.0.out_layers.0.weight': 'up_blocks.3.resnets.0.norm2.weight',
    'model.diffusion_model.output_blocks.9.0.out_layers.0.bias': 'up_blocks.3.resnets.0.norm2.bias',
    'model.diffusion_model.output_blocks.9.0.out_layers.3.weight': 'up_blocks.3.resnets.0.conv2.weight',
    'model.diffusion_model.output_blocks.9.0.out_layers.3.bias': 'up_blocks.3.resnets.0.conv2.bias',
    'model.diffusion_model.output_blocks.9.0.skip_connection.weight': 'up_blocks.3.resnets.0.conv_shortcut.weight',
    'model.diffusion_model.output_blocks.9.0.skip_connection.bias': 'up_blocks.3.resnets.0.conv_shortcut.bias',
    'model.diffusion_model.output_blocks.9.1.norm.weight': 'up_blocks.3.attentions.0.norm.weight',
    'model.diffusion_model.output_blocks.9.1.norm.bias': 'up_blocks.3.attentions.0.norm.bias',
    'model.diffusion_model.output_blocks.9.1.proj_in.weight': 'up_blocks.3.attentions.0.proj_in.weight',
    'model.diffusion_model.output_blocks.9.1.proj_in.bias': 'up_blocks.3.attentions.0.proj_in.bias',
    'model.diffusion_model.output_blocks.9.1.transformer_blocks.0.norm1.weight': 'up_blocks.3.attentions.0.transformer_blocks.0.norm1.weight',
    'model.diffusion_model.output_blocks.9.1.transformer_blocks.0.norm1.bias': 'up_blocks.3.attentions.0.transformer_blocks.0.norm1.bias',
    'model.diffusion_model.output_blocks.9.1.transformer_blocks.0.attn1.to_q.weight': 'up_blocks.3.attentions.0.transformer_blocks.0.attn1.to_q.weight',
    'model.diffusion_model.output_blocks.9.1.transformer_blocks.0.attn1.to_k.weight': 'up_blocks.3.attentions.0.transformer_blocks.0.attn1.to_k.weight',
    'model.diffusion_model.output_blocks.9.1.transformer_blocks.0.attn1.to_v.weight': 'up_blocks.3.attentions.0.transformer_blocks.0.attn1.to_v.weight',
    'model.diffusion_model.output_blocks.9.1.transformer_blocks.0.attn1.to_out.0.weight': 'up_blocks.3.attentions.0.transformer_blocks.0.attn1.to_out.0.weight',
    'model.diffusion_model.output_blocks.9.1.transformer_blocks.0.attn1.to_out.0.bias': 'up_blocks.3.attentions.0.transformer_blocks.0.attn1.to_out.0.bias',
    'model.diffusion_model.output_blocks.9.1.transformer_blocks.0.norm2.weight': 'up_blocks.3.attentions.0.transformer_blocks.0.norm2.weight',
    'model.diffusion_model.output_blocks.9.1.transformer_blocks.0.norm2.bias': 'up_blocks.3.attentions.0.transformer_blocks.0.norm2.bias',
    'model.diffusion_model.output_blocks.9.1.transformer_blocks.0.attn2.to_q.weight': 'up_blocks.3.attentions.0.transformer_blocks.0.attn2.to_q.weight',
    'model.diffusion_model.output_blocks.9.1.transformer_blocks.0.attn2.to_k.weight': 'up_blocks.3.attentions.0.transformer_blocks.0.attn2.to_k.weight',
    'model.diffusion_model.output_blocks.9.1.transformer_blocks.0.attn2.to_v.weight': 'up_blocks.3.attentions.0.transformer_blocks.0.attn2.to_v.weight',
    'model.diffusion_model.output_blocks.9.1.transformer_blocks.0.attn2.to_out.0.weight': 'up_blocks.3.attentions.0.transformer_blocks.0.attn2.to_out.0.weight',
    'model.diffusion_model.output_blocks.9.1.transformer_blocks.0.attn2.to_out.0.bias': 'up_blocks.3.attentions.0.transformer_blocks.0.attn2.to_out.0.bias',
    'model.diffusion_model.output_blocks.9.1.transformer_blocks.0.norm3.weight': 'up_blocks.3.attentions.0.transformer_blocks.0.norm3.weight',
    'model.diffusion_model.output_blocks.9.1.transformer_blocks.0.norm3.bias': 'up_blocks.3.attentions.0.transformer_blocks.0.norm3.bias',
    'model.diffusion_model.output_blocks.9.1.transformer_blocks.0.ff.net.0.proj.weight': 'up_blocks.3.attentions.0.transformer_blocks.0.ff.net.0.proj.weight',
    'model.diffusion_model.output_blocks.9.1.transformer_blocks.0.ff.net.0.proj.bias': 'up_blocks.3.attentions.0.transformer_blocks.0.ff.net.0.proj.bias',
    'model.diffusion_model.output_blocks.9.1.transformer_blocks.0.ff.net.2.weight': 'up_blocks.3.attentions.0.transformer_blocks.0.ff.net.2.weight',
    'model.diffusion_model.output_blocks.9.1.transformer_blocks.0.ff.net.2.bias': 'up_blocks.3.attentions.0.transformer_blocks.0.ff.net.2.bias',
    'model.diffusion_model.output_blocks.9.1.proj_out.weight': 'up_blocks.3.attentions.0.proj_out.weight',
    'model.diffusion_model.output_blocks.9.1.proj_out.bias': 'up_blocks.3.attentions.0.proj_out.bias',
    'model.diffusion_model.output_blocks.10.0.in_layers.0.weight': 'up_blocks.3.resnets.1.norm1.weight',
    'model.diffusion_model.output_blocks.10.0.in_layers.0.bias': 'up_blocks.3.resnets.1.norm1.bias',
    'model.diffusion_model.output_blocks.10.0.in_layers.2.weight': 'up_blocks.3.resnets.1.conv1.weight',
    'model.diffusion_model.output_blocks.10.0.in_layers.2.bias': 'up_blocks.3.resnets.1.conv1.bias',
    'model.diffusion_model.output_blocks.10.0.emb_layers.1.weight': 'up_blocks.3.resnets.1.time_emb_proj.weight',
    'model.diffusion_model.output_blocks.10.0.emb_layers.1.bias': 'up_blocks.3.resnets.1.time_emb_proj.bias',
    'model.diffusion_model.output_blocks.10.0.out_layers.0.weight': 'up_blocks.3.resnets.1.norm2.weight',
    'model.diffusion_model.output_blocks.10.0.out_layers.0.bias': 'up_blocks.3.resnets.1.norm2.bias',
    'model.diffusion_model.output_blocks.10.0.out_layers.3.weight': 'up_blocks.3.resnets.1.conv2.weight',
    'model.diffusion_model.output_blocks.10.0.out_layers.3.bias': 'up_blocks.3.resnets.1.conv2.bias',
    'model.diffusion_model.output_blocks.10.0.skip_connection.weight': 'up_blocks.3.resnets.1.conv_shortcut.weight',
    'model.diffusion_model.output_blocks.10.0.skip_connection.bias': 'up_blocks.3.resnets.1.conv_shortcut.bias',
    'model.diffusion_model.output_blocks.10.1.norm.weight': 'up_blocks.3.attentions.1.norm.weight',
    'model.diffusion_model.output_blocks.10.1.norm.bias': 'up_blocks.3.attentions.1.norm.bias',
    'model.diffusion_model.output_blocks.10.1.proj_in.weight': 'up_blocks.3.attentions.1.proj_in.weight',
    'model.diffusion_model.output_blocks.10.1.proj_in.bias': 'up_blocks.3.attentions.1.proj_in.bias',
    'model.diffusion_model.output_blocks.10.1.transformer_blocks.0.norm1.weight': 'up_blocks.3.attentions.1.transformer_blocks.0.norm1.weight',
    'model.diffusion_model.output_blocks.10.1.transformer_blocks.0.norm1.bias': 'up_blocks.3.attentions.1.transformer_blocks.0.norm1.bias',
    'model.diffusion_model.output_blocks.10.1.transformer_blocks.0.attn1.to_q.weight': 'up_blocks.3.attentions.1.transformer_blocks.0.attn1.to_q.weight',
    'model.diffusion_model.output_blocks.10.1.transformer_blocks.0.attn1.to_k.weight': 'up_blocks.3.attentions.1.transformer_blocks.0.attn1.to_k.weight',
    'model.diffusion_model.output_blocks.10.1.transformer_blocks.0.attn1.to_v.weight': 'up_blocks.3.attentions.1.transformer_blocks.0.attn1.to_v.weight',
    'model.diffusion_model.output_blocks.10.1.transformer_blocks.0.attn1.to_out.0.weight': 'up_blocks.3.attentions.1.transformer_blocks.0.attn1.to_out.0.weight',
    'model.diffusion_model.output_blocks.10.1.transformer_blocks.0.attn1.to_out.0.bias': 'up_blocks.3.attentions.1.transformer_blocks.0.attn1.to_out.0.bias',
    'model.diffusion_model.output_blocks.10.1.transformer_blocks.0.norm2.weight': 'up_blocks.3.attentions.1.transformer_blocks.0.norm2.weight',
    'model.diffusion_model.output_blocks.10.1.transformer_blocks.0.norm2.bias': 'up_blocks.3.attentions.1.transformer_blocks.0.norm2.bias',
    'model.diffusion_model.output_blocks.10.1.transformer_blocks.0.attn2.to_q.weight': 'up_blocks.3.attentions.1.transformer_blocks.0.attn2.to_q.weight',
    'model.diffusion_model.output_blocks.10.1.transformer_blocks.0.attn2.to_k.weight': 'up_blocks.3.attentions.1.transformer_blocks.0.attn2.to_k.weight',
    'model.diffusion_model.output_blocks.10.1.transformer_blocks.0.attn2.to_v.weight': 'up_blocks.3.attentions.1.transformer_blocks.0.attn2.to_v.weight',
    'model.diffusion_model.output_blocks.10.1.transformer_blocks.0.attn2.to_out.0.weight': 'up_blocks.3.attentions.1.transformer_blocks.0.attn2.to_out.0.weight',
    'model.diffusion_model.output_blocks.10.1.transformer_blocks.0.attn2.to_out.0.bias': 'up_blocks.3.attentions.1.transformer_blocks.0.attn2.to_out.0.bias',
    'model.diffusion_model.output_blocks.10.1.transformer_blocks.0.norm3.weight': 'up_blocks.3.attentions.1.transformer_blocks.0.norm3.weight',
    'model.diffusion_model.output_blocks.10.1.transformer_blocks.0.norm3.bias': 'up_blocks.3.attentions.1.transformer_blocks.0.norm3.bias',
    'model.diffusion_model.output_blocks.10.1.transformer_blocks.0.ff.net.0.proj.weight': 'up_blocks.3.attentions.1.transformer_blocks.0.ff.net.0.proj.weight',
    'model.diffusion_model.output_blocks.10.1.transformer_blocks.0.ff.net.0.proj.bias': 'up_blocks.3.attentions.1.transformer_blocks.0.ff.net.0.proj.bias',
    'model.diffusion_model.output_blocks.10.1.transformer_blocks.0.ff.net.2.weight': 'up_blocks.3.attentions.1.transformer_blocks.0.ff.net.2.weight',
    'model.diffusion_model.output_blocks.10.1.transformer_blocks.0.ff.net.2.bias': 'up_blocks.3.attentions.1.transformer_blocks.0.ff.net.2.bias',
    'model.diffusion_model.output_blocks.10.1.proj_out.weight': 'up_blocks.3.attentions.1.proj_out.weight',
    'model.diffusion_model.output_blocks.10.1.proj_out.bias': 'up_blocks.3.attentions.1.proj_out.bias',
    'model.diffusion_model.output_blocks.11.0.in_layers.0.weight': 'up_blocks.3.resnets.2.norm1.weight',
    'model.diffusion_model.output_blocks.11.0.in_layers.0.bias': 'up_blocks.3.resnets.2.norm1.bias',
    'model.diffusion_model.output_blocks.11.0.in_layers.2.weight': 'up_blocks.3.resnets.2.conv1.weight',
    'model.diffusion_model.output_blocks.11.0.in_layers.2.bias': 'up_blocks.3.resnets.2.conv1.bias',
    'model.diffusion_model.output_blocks.11.0.emb_layers.1.weight': 'up_blocks.3.resnets.2.time_emb_proj.weight',
    'model.diffusion_model.output_blocks.11.0.emb_layers.1.bias': 'up_blocks.3.resnets.2.time_emb_proj.bias',
    'model.diffusion_model.output_blocks.11.0.out_layers.0.weight': 'up_blocks.3.resnets.2.norm2.weight',
    'model.diffusion_model.output_blocks.11.0.out_layers.0.bias': 'up_blocks.3.resnets.2.norm2.bias',
    'model.diffusion_model.output_blocks.11.0.out_layers.3.weight': 'up_blocks.3.resnets.2.conv2.weight',
    'model.diffusion_model.output_blocks.11.0.out_layers.3.bias': 'up_blocks.3.resnets.2.conv2.bias',
    'model.diffusion_model.output_blocks.11.0.skip_connection.weight': 'up_blocks.3.resnets.2.conv_shortcut.weight',
    'model.diffusion_model.output_blocks.11.0.skip_connection.bias': 'up_blocks.3.resnets.2.conv_shortcut.bias',
    'model.diffusion_model.output_blocks.11.1.norm.weight': 'up_blocks.3.attentions.2.norm.weight',
    'model.diffusion_model.output_blocks.11.1.norm.bias': 'up_blocks.3.attentions.2.norm.bias',
    'model.diffusion_model.output_blocks.11.1.proj_in.weight': 'up_blocks.3.attentions.2.proj_in.weight',
    'model.diffusion_model.output_blocks.11.1.proj_in.bias': 'up_blocks.3.attentions.2.proj_in.bias',
    'model.diffusion_model.output_blocks.11.1.transformer_blocks.0.norm1.weight': 'up_blocks.3.attentions.2.transformer_blocks.0.norm1.weight',
    'model.diffusion_model.output_blocks.11.1.transformer_blocks.0.norm1.bias': 'up_blocks.3.attentions.2.transformer_blocks.0.norm1.bias',
    'model.diffusion_model.output_blocks.11.1.transformer_blocks.0.attn1.to_q.weight': 'up_blocks.3.attentions.2.transformer_blocks.0.attn1.to_q.weight',
    'model.diffusion_model.output_blocks.11.1.transformer_blocks.0.attn1.to_k.weight': 'up_blocks.3.attentions.2.transformer_blocks.0.attn1.to_k.weight',
    'model.diffusion_model.output_blocks.11.1.transformer_blocks.0.attn1.to_v.weight': 'up_blocks.3.attentions.2.transformer_blocks.0.attn1.to_v.weight',
    'model.diffusion_model.output_blocks.11.1.transformer_blocks.0.attn1.to_out.0.weight': 'up_blocks.3.attentions.2.transformer_blocks.0.attn1.to_out.0.weight',
    'model.diffusion_model.output_blocks.11.1.transformer_blocks.0.attn1.to_out.0.bias': 'up_blocks.3.attentions.2.transformer_blocks.0.attn1.to_out.0.bias',
    'model.diffusion_model.output_blocks.11.1.transformer_blocks.0.norm2.weight': 'up_blocks.3.attentions.2.transformer_blocks.0.norm2.weight',
    'model.diffusion_model.output_blocks.11.1.transformer_blocks.0.norm2.bias': 'up_blocks.3.attentions.2.transformer_blocks.0.norm2.bias',
    'model.diffusion_model.output_blocks.11.1.transformer_blocks.0.attn2.to_q.weight': 'up_blocks.3.attentions.2.transformer_blocks.0.attn2.to_q.weight',
    'model.diffusion_model.output_blocks.11.1.transformer_blocks.0.attn2.to_k.weight': 'up_blocks.3.attentions.2.transformer_blocks.0.attn2.to_k.weight',
    'model.diffusion_model.output_blocks.11.1.transformer_blocks.0.attn2.to_v.weight': 'up_blocks.3.attentions.2.transformer_blocks.0.attn2.to_v.weight',
    'model.diffusion_model.output_blocks.11.1.transformer_blocks.0.attn2.to_out.0.weight': 'up_blocks.3.attentions.2.transformer_blocks.0.attn2.to_out.0.weight',
    'model.diffusion_model.output_blocks.11.1.transformer_blocks.0.attn2.to_out.0.bias': 'up_blocks.3.attentions.2.transformer_blocks.0.attn2.to_out.0.bias',
    'model.diffusion_model.output_blocks.11.1.transformer_blocks.0.norm3.weight': 'up_blocks.3.attentions.2.transformer_blocks.0.norm3.weight',
    'model.diffusion_model.output_blocks.11.1.transformer_blocks.0.norm3.bias': 'up_blocks.3.attentions.2.transformer_blocks.0.norm3.bias',
    'model.diffusion_model.output_blocks.11.1.transformer_blocks.0.ff.net.0.proj.weight': 'up_blocks.3.attentions.2.transformer_blocks.0.ff.net.0.proj.weight',
    'model.diffusion_model.output_blocks.11.1.transformer_blocks.0.ff.net.0.proj.bias': 'up_blocks.3.attentions.2.transformer_blocks.0.ff.net.0.proj.bias',
    'model.diffusion_model.output_blocks.11.1.transformer_blocks.0.ff.net.2.weight': 'up_blocks.3.attentions.2.transformer_blocks.0.ff.net.2.weight',
    'model.diffusion_model.output_blocks.11.1.transformer_blocks.0.ff.net.2.bias': 'up_blocks.3.attentions.2.transformer_blocks.0.ff.net.2.bias',
    'model.diffusion_model.output_blocks.11.1.proj_out.weight': 'up_blocks.3.attentions.2.proj_out.weight',
    'model.diffusion_model.output_blocks.11.1.proj_out.bias': 'up_blocks.3.attentions.2.proj_out.bias',
    'model.diffusion_model.out.0.weight': 'conv_norm_out.weight',
    'model.diffusion_model.out.0.bias': 'conv_norm_out.bias',
    'model.diffusion_model.out.2.weight': 'conv_out.weight',
    'model.diffusion_model.out.2.bias': 'conv_out.bias',
}
CKPT_MAPPING = {
    'civitai_model': [('model.diffusion_model.time_embed.0.weight', (1, 0)),
                      ('model.diffusion_model.time_embed.0.bias', None),
                      ('model.diffusion_model.input_blocks.0.0.weight', (2, 3, 1, 0)),
                      ('model.diffusion_model.input_blocks.0.0.bias', None),
                      ('model.diffusion_model.time_embed.2.weight', (1, 0)),
                      ('model.diffusion_model.time_embed.2.bias', None),
                      ('model.diffusion_model.input_blocks.1.0.in_layers.0.weight', None),
                      ('model.diffusion_model.input_blocks.1.0.in_layers.0.bias', None),
                      ('model.diffusion_model.input_blocks.1.0.in_layers.2.weight', (2, 3, 1, 0)),
                      ('model.diffusion_model.input_blocks.1.0.in_layers.2.bias', None),
                      ('model.diffusion_model.input_blocks.1.0.emb_layers.1.weight', (1, 0)),
                      ('model.diffusion_model.input_blocks.1.0.emb_layers.1.bias', None),
                      ('model.diffusion_model.input_blocks.1.0.out_layers.0.weight', None),
                      ('model.diffusion_model.input_blocks.1.0.out_layers.0.bias', None),
                      ('model.diffusion_model.input_blocks.1.0.out_layers.3.weight', (2, 3, 1, 0)),
                      ('model.diffusion_model.input_blocks.1.0.out_layers.3.bias', None),
                      ('model.diffusion_model.input_blocks.1.1.norm.weight', None),
                      ('model.diffusion_model.input_blocks.1.1.norm.bias', None),
                      ('model.diffusion_model.input_blocks.1.1.proj_in.weight', (2, 3, 1, 0)),
                      ('model.diffusion_model.input_blocks.1.1.proj_in.bias', None),
                      ('model.diffusion_model.input_blocks.1.1.transformer_blocks.0.norm1.weight', None),
                      ('model.diffusion_model.input_blocks.1.1.transformer_blocks.0.norm1.bias', None),
                      ('model.diffusion_model.input_blocks.1.1.transformer_blocks.0.attn1.to_q.weight', (1, 0)),
                      ('model.diffusion_model.input_blocks.1.1.transformer_blocks.0.attn1.to_k.weight', (1, 0)),
                      ('model.diffusion_model.input_blocks.1.1.transformer_blocks.0.attn1.to_v.weight', (1, 0)),
                      ('model.diffusion_model.input_blocks.1.1.transformer_blocks.0.attn1.to_out.0.weight', (1, 0)),
                      ('model.diffusion_model.input_blocks.1.1.transformer_blocks.0.attn1.to_out.0.bias', None),
                      ('model.diffusion_model.input_blocks.1.1.transformer_blocks.0.norm2.weight', None),
                      ('model.diffusion_model.input_blocks.1.1.transformer_blocks.0.norm2.bias', None),
                      ('model.diffusion_model.input_blocks.1.1.transformer_blocks.0.attn2.to_q.weight', (1, 0)),
                      ('model.diffusion_model.input_blocks.1.1.transformer_blocks.0.attn2.to_k.weight', (1, 0)),
                      ('model.diffusion_model.input_blocks.1.1.transformer_blocks.0.attn2.to_v.weight', (1, 0)),
                      ('model.diffusion_model.input_blocks.1.1.transformer_blocks.0.attn2.to_out.0.weight', (1, 0)),
                      ('model.diffusion_model.input_blocks.1.1.transformer_blocks.0.attn2.to_out.0.bias', None),
                      ('model.diffusion_model.input_blocks.1.1.transformer_blocks.0.norm3.weight', None),
                      ('model.diffusion_model.input_blocks.1.1.transformer_blocks.0.norm3.bias', None),
                      ('model.diffusion_model.input_blocks.1.1.transformer_blocks.0.ff.net.0.proj.weight', (1, 0)),
                      ('model.diffusion_model.input_blocks.1.1.transformer_blocks.0.ff.net.0.proj.bias', None),
                      ('model.diffusion_model.input_blocks.1.1.transformer_blocks.0.ff.net.2.weight', (1, 0)),
                      ('model.diffusion_model.input_blocks.1.1.transformer_blocks.0.ff.net.2.bias', None),
                      ('model.diffusion_model.input_blocks.1.1.proj_out.weight', (2, 3, 1, 0)),
                      ('model.diffusion_model.input_blocks.1.1.proj_out.bias', None),
                      ('model.diffusion_model.input_blocks.2.0.in_layers.0.weight', None),
                      ('model.diffusion_model.input_blocks.2.0.in_layers.0.bias', None),
                      ('model.diffusion_model.input_blocks.2.0.in_layers.2.weight', (2, 3, 1, 0)),
                      ('model.diffusion_model.input_blocks.2.0.in_layers.2.bias', None),
                      ('model.diffusion_model.input_blocks.2.0.emb_layers.1.weight', (1, 0)),
                      ('model.diffusion_model.input_blocks.2.0.emb_layers.1.bias', None),
                      ('model.diffusion_model.input_blocks.2.0.out_layers.0.weight', None),
                      ('model.diffusion_model.input_blocks.2.0.out_layers.0.bias', None),
                      ('model.diffusion_model.input_blocks.2.0.out_layers.3.weight', (2, 3, 1, 0)),
                      ('model.diffusion_model.input_blocks.2.0.out_layers.3.bias', None),
                      ('model.diffusion_model.input_blocks.2.1.norm.weight', None),
                      ('model.diffusion_model.input_blocks.2.1.norm.bias', None),
                      ('model.diffusion_model.input_blocks.2.1.proj_in.weight', (2, 3, 1, 0)),
                      ('model.diffusion_model.input_blocks.2.1.proj_in.bias', None),
                      ('model.diffusion_model.input_blocks.2.1.transformer_blocks.0.norm1.weight', None),
                      ('model.diffusion_model.input_blocks.2.1.transformer_blocks.0.norm1.bias', None),
                      ('model.diffusion_model.input_blocks.2.1.transformer_blocks.0.attn1.to_q.weight', (1, 0)),
                      ('model.diffusion_model.input_blocks.2.1.transformer_blocks.0.attn1.to_k.weight', (1, 0)),
                      ('model.diffusion_model.input_blocks.2.1.transformer_blocks.0.attn1.to_v.weight', (1, 0)),
                      ('model.diffusion_model.input_blocks.2.1.transformer_blocks.0.attn1.to_out.0.weight', (1, 0)),
                      ('model.diffusion_model.input_blocks.2.1.transformer_blocks.0.attn1.to_out.0.bias', None),
                      ('model.diffusion_model.input_blocks.2.1.transformer_blocks.0.norm2.weight', None),
                      ('model.diffusion_model.input_blocks.2.1.transformer_blocks.0.norm2.bias', None),
                      ('model.diffusion_model.input_blocks.2.1.transformer_blocks.0.attn2.to_q.weight', (1, 0)),
                      ('model.diffusion_model.input_blocks.2.1.transformer_blocks.0.attn2.to_k.weight', (1, 0)),
                      ('model.diffusion_model.input_blocks.2.1.transformer_blocks.0.attn2.to_v.weight', (1, 0)),
                      ('model.diffusion_model.input_blocks.2.1.transformer_blocks.0.attn2.to_out.0.weight', (1, 0)),
                      ('model.diffusion_model.input_blocks.2.1.transformer_blocks.0.attn2.to_out.0.bias', None),
                      ('model.diffusion_model.input_blocks.2.1.transformer_blocks.0.norm3.weight', None),
                      ('model.diffusion_model.input_blocks.2.1.transformer_blocks.0.norm3.bias', None),
                      ('model.diffusion_model.input_blocks.2.1.transformer_blocks.0.ff.net.0.proj.weight', (1, 0)),
                      ('model.diffusion_model.input_blocks.2.1.transformer_blocks.0.ff.net.0.proj.bias', None),
                      ('model.diffusion_model.input_blocks.2.1.transformer_blocks.0.ff.net.2.weight', (1, 0)),
                      ('model.diffusion_model.input_blocks.2.1.transformer_blocks.0.ff.net.2.bias', None),
                      ('model.diffusion_model.input_blocks.2.1.proj_out.weight', (2, 3, 1, 0)),
                      ('model.diffusion_model.input_blocks.2.1.proj_out.bias', None),
                      ('model.diffusion_model.input_blocks.3.0.op.weight', (2, 3, 1, 0)),
                      ('model.diffusion_model.input_blocks.3.0.op.bias', None),
                      ('model.diffusion_model.input_blocks.4.0.in_layers.0.weight', None),
                      ('model.diffusion_model.input_blocks.4.0.in_layers.0.bias', None),
                      ('model.diffusion_model.input_blocks.4.0.in_layers.2.weight', (2, 3, 1, 0)),
                      ('model.diffusion_model.input_blocks.4.0.in_layers.2.bias', None),
                      ('model.diffusion_model.input_blocks.4.0.emb_layers.1.weight', (1, 0)),
                      ('model.diffusion_model.input_blocks.4.0.emb_layers.1.bias', None),
                      ('model.diffusion_model.input_blocks.4.0.out_layers.0.weight', None),
                      ('model.diffusion_model.input_blocks.4.0.out_layers.0.bias', None),
                      ('model.diffusion_model.input_blocks.4.0.out_layers.3.weight', (2, 3, 1, 0)),
                      ('model.diffusion_model.input_blocks.4.0.out_layers.3.bias', None),
                      ('model.diffusion_model.input_blocks.4.0.skip_connection.weight', (2, 3, 1, 0)),
                      ('model.diffusion_model.input_blocks.4.0.skip_connection.bias', None),
                      ('model.diffusion_model.input_blocks.4.1.norm.weight', None),
                      ('model.diffusion_model.input_blocks.4.1.norm.bias', None),
                      ('model.diffusion_model.input_blocks.4.1.proj_in.weight', (2, 3, 1, 0)),
                      ('model.diffusion_model.input_blocks.4.1.proj_in.bias', None),
                      ('model.diffusion_model.input_blocks.4.1.transformer_blocks.0.norm1.weight', None),
                      ('model.diffusion_model.input_blocks.4.1.transformer_blocks.0.norm1.bias', None),
                      ('model.diffusion_model.input_blocks.4.1.transformer_blocks.0.attn1.to_q.weight', (1, 0)),
                      ('model.diffusion_model.input_blocks.4.1.transformer_blocks.0.attn1.to_k.weight', (1, 0)),
                      ('model.diffusion_model.input_blocks.4.1.transformer_blocks.0.attn1.to_v.weight', (1, 0)),
                      ('model.diffusion_model.input_blocks.4.1.transformer_blocks.0.attn1.to_out.0.weight', (1, 0)),
                      ('model.diffusion_model.input_blocks.4.1.transformer_blocks.0.attn1.to_out.0.bias', None),
                      ('model.diffusion_model.input_blocks.4.1.transformer_blocks.0.norm2.weight', None),
                      ('model.diffusion_model.input_blocks.4.1.transformer_blocks.0.norm2.bias', None),
                      ('model.diffusion_model.input_blocks.4.1.transformer_blocks.0.attn2.to_q.weight', (1, 0)),
                      ('model.diffusion_model.input_blocks.4.1.transformer_blocks.0.attn2.to_k.weight', (1, 0)),
                      ('model.diffusion_model.input_blocks.4.1.transformer_blocks.0.attn2.to_v.weight', (1, 0)),
                      ('model.diffusion_model.input_blocks.4.1.transformer_blocks.0.attn2.to_out.0.weight', (1, 0)),
                      ('model.diffusion_model.input_blocks.4.1.transformer_blocks.0.attn2.to_out.0.bias', None),
                      ('model.diffusion_model.input_blocks.4.1.transformer_blocks.0.norm3.weight', None),
                      ('model.diffusion_model.input_blocks.4.1.transformer_blocks.0.norm3.bias', None),
                      ('model.diffusion_model.input_blocks.4.1.transformer_blocks.0.ff.net.0.proj.weight', (1, 0)),
                      ('model.diffusion_model.input_blocks.4.1.transformer_blocks.0.ff.net.0.proj.bias', None),
                      ('model.diffusion_model.input_blocks.4.1.transformer_blocks.0.ff.net.2.weight', (1, 0)),
                      ('model.diffusion_model.input_blocks.4.1.transformer_blocks.0.ff.net.2.bias', None),
                      ('model.diffusion_model.input_blocks.4.1.proj_out.weight', (2, 3, 1, 0)),
                      ('model.diffusion_model.input_blocks.4.1.proj_out.bias', None),
                      ('model.diffusion_model.input_blocks.5.0.in_layers.0.weight', None),
                      ('model.diffusion_model.input_blocks.5.0.in_layers.0.bias', None),
                      ('model.diffusion_model.input_blocks.5.0.in_layers.2.weight', (2, 3, 1, 0)),
                      ('model.diffusion_model.input_blocks.5.0.in_layers.2.bias', None),
                      ('model.diffusion_model.input_blocks.5.0.emb_layers.1.weight', (1, 0)),
                      ('model.diffusion_model.input_blocks.5.0.emb_layers.1.bias', None),
                      ('model.diffusion_model.input_blocks.5.0.out_layers.0.weight', None),
                      ('model.diffusion_model.input_blocks.5.0.out_layers.0.bias', None),
                      ('model.diffusion_model.input_blocks.5.0.out_layers.3.weight', (2, 3, 1, 0)),
                      ('model.diffusion_model.input_blocks.5.0.out_layers.3.bias', None),
                      ('model.diffusion_model.input_blocks.5.1.norm.weight', None),
                      ('model.diffusion_model.input_blocks.5.1.norm.bias', None),
                      ('model.diffusion_model.input_blocks.5.1.proj_in.weight', (2, 3, 1, 0)),
                      ('model.diffusion_model.input_blocks.5.1.proj_in.bias', None),
                      ('model.diffusion_model.input_blocks.5.1.transformer_blocks.0.norm1.weight', None),
                      ('model.diffusion_model.input_blocks.5.1.transformer_blocks.0.norm1.bias', None),
                      ('model.diffusion_model.input_blocks.5.1.transformer_blocks.0.attn1.to_q.weight', (1, 0)),
                      ('model.diffusion_model.input_blocks.5.1.transformer_blocks.0.attn1.to_k.weight', (1, 0)),
                      ('model.diffusion_model.input_blocks.5.1.transformer_blocks.0.attn1.to_v.weight', (1, 0)),
                      ('model.diffusion_model.input_blocks.5.1.transformer_blocks.0.attn1.to_out.0.weight', (1, 0)),
                      ('model.diffusion_model.input_blocks.5.1.transformer_blocks.0.attn1.to_out.0.bias', None),
                      ('model.diffusion_model.input_blocks.5.1.transformer_blocks.0.norm2.weight', None),
                      ('model.diffusion_model.input_blocks.5.1.transformer_blocks.0.norm2.bias', None),
                      ('model.diffusion_model.input_blocks.5.1.transformer_blocks.0.attn2.to_q.weight', (1, 0)),
                      ('model.diffusion_model.input_blocks.5.1.transformer_blocks.0.attn2.to_k.weight', (1, 0)),
                      ('model.diffusion_model.input_blocks.5.1.transformer_blocks.0.attn2.to_v.weight', (1, 0)),
                      ('model.diffusion_model.input_blocks.5.1.transformer_blocks.0.attn2.to_out.0.weight', (1, 0)),
                      ('model.diffusion_model.input_blocks.5.1.transformer_blocks.0.attn2.to_out.0.bias', None),
                      ('model.diffusion_model.input_blocks.5.1.transformer_blocks.0.norm3.weight', None),
                      ('model.diffusion_model.input_blocks.5.1.transformer_blocks.0.norm3.bias', None),
                      ('model.diffusion_model.input_blocks.5.1.transformer_blocks.0.ff.net.0.proj.weight', (1, 0)),
                      ('model.diffusion_model.input_blocks.5.1.transformer_blocks.0.ff.net.0.proj.bias', None),
                      ('model.diffusion_model.input_blocks.5.1.transformer_blocks.0.ff.net.2.weight', (1, 0)),
                      ('model.diffusion_model.input_blocks.5.1.transformer_blocks.0.ff.net.2.bias', None),
                      ('model.diffusion_model.input_blocks.5.1.proj_out.weight', (2, 3, 1, 0)),
                      ('model.diffusion_model.input_blocks.5.1.proj_out.bias', None),
                      ('model.diffusion_model.input_blocks.6.0.op.weight', (2, 3, 1, 0)),
                      ('model.diffusion_model.input_blocks.6.0.op.bias', None),
                      ('model.diffusion_model.input_blocks.7.0.in_layers.0.weight', None),
                      ('model.diffusion_model.input_blocks.7.0.in_layers.0.bias', None),
                      ('model.diffusion_model.input_blocks.7.0.in_layers.2.weight', (2, 3, 1, 0)),
                      ('model.diffusion_model.input_blocks.7.0.in_layers.2.bias', None),
                      ('model.diffusion_model.input_blocks.7.0.emb_layers.1.weight', (1, 0)),
                      ('model.diffusion_model.input_blocks.7.0.emb_layers.1.bias', None),
                      ('model.diffusion_model.input_blocks.7.0.out_layers.0.weight', None),
                      ('model.diffusion_model.input_blocks.7.0.out_layers.0.bias', None),
                      ('model.diffusion_model.input_blocks.7.0.out_layers.3.weight', (2, 3, 1, 0)),
                      ('model.diffusion_model.input_blocks.7.0.out_layers.3.bias', None),
                      ('model.diffusion_model.input_blocks.7.0.skip_connection.weight', (2, 3, 1, 0)),
                      ('model.diffusion_model.input_blocks.7.0.skip_connection.bias', None),
                      ('model.diffusion_model.input_blocks.7.1.norm.weight', None),
                      ('model.diffusion_model.input_blocks.7.1.norm.bias', None),
                      ('model.diffusion_model.input_blocks.7.1.proj_in.weight', (2, 3, 1, 0)),
                      ('model.diffusion_model.input_blocks.7.1.proj_in.bias', None),
                      ('model.diffusion_model.input_blocks.7.1.transformer_blocks.0.norm1.weight', None),
                      ('model.diffusion_model.input_blocks.7.1.transformer_blocks.0.norm1.bias', None),
                      ('model.diffusion_model.input_blocks.7.1.transformer_blocks.0.attn1.to_q.weight', (1, 0)),
                      ('model.diffusion_model.input_blocks.7.1.transformer_blocks.0.attn1.to_k.weight', (1, 0)),
                      ('model.diffusion_model.input_blocks.7.1.transformer_blocks.0.attn1.to_v.weight', (1, 0)),
                      ('model.diffusion_model.input_blocks.7.1.transformer_blocks.0.attn1.to_out.0.weight', (1, 0)),
                      ('model.diffusion_model.input_blocks.7.1.transformer_blocks.0.attn1.to_out.0.bias', None),
                      ('model.diffusion_model.input_blocks.7.1.transformer_blocks.0.norm2.weight', None),
                      ('model.diffusion_model.input_blocks.7.1.transformer_blocks.0.norm2.bias', None),
                      ('model.diffusion_model.input_blocks.7.1.transformer_blocks.0.attn2.to_q.weight', (1, 0)),
                      ('model.diffusion_model.input_blocks.7.1.transformer_blocks.0.attn2.to_k.weight', (1, 0)),
                      ('model.diffusion_model.input_blocks.7.1.transformer_blocks.0.attn2.to_v.weight', (1, 0)),
                      ('model.diffusion_model.input_blocks.7.1.transformer_blocks.0.attn2.to_out.0.weight', (1, 0)),
                      ('model.diffusion_model.input_blocks.7.1.transformer_blocks.0.attn2.to_out.0.bias', None),
                      ('model.diffusion_model.input_blocks.7.1.transformer_blocks.0.norm3.weight', None),
                      ('model.diffusion_model.input_blocks.7.1.transformer_blocks.0.norm3.bias', None),
                      ('model.diffusion_model.input_blocks.7.1.transformer_blocks.0.ff.net.0.proj.weight', (1, 0)),
                      ('model.diffusion_model.input_blocks.7.1.transformer_blocks.0.ff.net.0.proj.bias', None),
                      ('model.diffusion_model.input_blocks.7.1.transformer_blocks.0.ff.net.2.weight', (1, 0)),
                      ('model.diffusion_model.input_blocks.7.1.transformer_blocks.0.ff.net.2.bias', None),
                      ('model.diffusion_model.input_blocks.7.1.proj_out.weight', (2, 3, 1, 0)),
                      ('model.diffusion_model.input_blocks.7.1.proj_out.bias', None),
                      ('model.diffusion_model.input_blocks.8.0.in_layers.0.weight', None),
                      ('model.diffusion_model.input_blocks.8.0.in_layers.0.bias', None),
                      ('model.diffusion_model.input_blocks.8.0.in_layers.2.weight', (2, 3, 1, 0)),
                      ('model.diffusion_model.input_blocks.8.0.in_layers.2.bias', None),
                      ('model.diffusion_model.input_blocks.8.0.emb_layers.1.weight', (1, 0)),
                      ('model.diffusion_model.input_blocks.8.0.emb_layers.1.bias', None),
                      ('model.diffusion_model.input_blocks.8.0.out_layers.0.weight', None),
                      ('model.diffusion_model.input_blocks.8.0.out_layers.0.bias', None),
                      ('model.diffusion_model.input_blocks.8.0.out_layers.3.weight', (2, 3, 1, 0)),
                      ('model.diffusion_model.input_blocks.8.0.out_layers.3.bias', None),
                      ('model.diffusion_model.input_blocks.8.1.norm.weight', None),
                      ('model.diffusion_model.input_blocks.8.1.norm.bias', None),
                      ('model.diffusion_model.input_blocks.8.1.proj_in.weight', (2, 3, 1, 0)),
                      ('model.diffusion_model.input_blocks.8.1.proj_in.bias', None),
                      ('model.diffusion_model.input_blocks.8.1.transformer_blocks.0.norm1.weight', None),
                      ('model.diffusion_model.input_blocks.8.1.transformer_blocks.0.norm1.bias', None),
                      ('model.diffusion_model.input_blocks.8.1.transformer_blocks.0.attn1.to_q.weight', (1, 0)),
                      ('model.diffusion_model.input_blocks.8.1.transformer_blocks.0.attn1.to_k.weight', (1, 0)),
                      ('model.diffusion_model.input_blocks.8.1.transformer_blocks.0.attn1.to_v.weight', (1, 0)),
                      ('model.diffusion_model.input_blocks.8.1.transformer_blocks.0.attn1.to_out.0.weight', (1, 0)),
                      ('model.diffusion_model.input_blocks.8.1.transformer_blocks.0.attn1.to_out.0.bias', None),
                      ('model.diffusion_model.input_blocks.8.1.transformer_blocks.0.norm2.weight', None),
                      ('model.diffusion_model.input_blocks.8.1.transformer_blocks.0.norm2.bias', None),
                      ('model.diffusion_model.input_blocks.8.1.transformer_blocks.0.attn2.to_q.weight', (1, 0)),
                      ('model.diffusion_model.input_blocks.8.1.transformer_blocks.0.attn2.to_k.weight', (1, 0)),
                      ('model.diffusion_model.input_blocks.8.1.transformer_blocks.0.attn2.to_v.weight', (1, 0)),
                      ('model.diffusion_model.input_blocks.8.1.transformer_blocks.0.attn2.to_out.0.weight', (1, 0)),
                      ('model.diffusion_model.input_blocks.8.1.transformer_blocks.0.attn2.to_out.0.bias', None),
                      ('model.diffusion_model.input_blocks.8.1.transformer_blocks.0.norm3.weight', None),
                      ('model.diffusion_model.input_blocks.8.1.transformer_blocks.0.norm3.bias', None),
                      ('model.diffusion_model.input_blocks.8.1.transformer_blocks.0.ff.net.0.proj.weight', (1, 0)),
                      ('model.diffusion_model.input_blocks.8.1.transformer_blocks.0.ff.net.0.proj.bias', None),
                      ('model.diffusion_model.input_blocks.8.1.transformer_blocks.0.ff.net.2.weight', (1, 0)),
                      ('model.diffusion_model.input_blocks.8.1.transformer_blocks.0.ff.net.2.bias', None),
                      ('model.diffusion_model.input_blocks.8.1.proj_out.weight', (2, 3, 1, 0)),
                      ('model.diffusion_model.input_blocks.8.1.proj_out.bias', None),
                      ('model.diffusion_model.input_blocks.9.0.op.weight', (2, 3, 1, 0)),
                      ('model.diffusion_model.input_blocks.9.0.op.bias', None),
                      ('model.diffusion_model.input_blocks.10.0.in_layers.0.weight', None),
                      ('model.diffusion_model.input_blocks.10.0.in_layers.0.bias', None),
                      ('model.diffusion_model.input_blocks.10.0.in_layers.2.weight', (2, 3, 1, 0)),
                      ('model.diffusion_model.input_blocks.10.0.in_layers.2.bias', None),
                      ('model.diffusion_model.input_blocks.10.0.emb_layers.1.weight', (1, 0)),
                      ('model.diffusion_model.input_blocks.10.0.emb_layers.1.bias', None),
                      ('model.diffusion_model.input_blocks.10.0.out_layers.0.weight', None),
                      ('model.diffusion_model.input_blocks.10.0.out_layers.0.bias', None),
                      ('model.diffusion_model.input_blocks.10.0.out_layers.3.weight', (2, 3, 1, 0)),
                      ('model.diffusion_model.input_blocks.10.0.out_layers.3.bias', None),
                      ('model.diffusion_model.input_blocks.11.0.in_layers.0.weight', None),
                      ('model.diffusion_model.input_blocks.11.0.in_layers.0.bias', None),
                      ('model.diffusion_model.input_blocks.11.0.in_layers.2.weight', (2, 3, 1, 0)),
                      ('model.diffusion_model.input_blocks.11.0.in_layers.2.bias', None),
                      ('model.diffusion_model.input_blocks.11.0.emb_layers.1.weight', (1, 0)),
                      ('model.diffusion_model.input_blocks.11.0.emb_layers.1.bias', None),
                      ('model.diffusion_model.input_blocks.11.0.out_layers.0.weight', None),
                      ('model.diffusion_model.input_blocks.11.0.out_layers.0.bias', None),
                      ('model.diffusion_model.input_blocks.11.0.out_layers.3.weight', (2, 3, 1, 0)),
                      ('model.diffusion_model.input_blocks.11.0.out_layers.3.bias', None),
                      ('model.diffusion_model.middle_block.0.in_layers.0.weight', None),
                      ('model.diffusion_model.middle_block.0.in_layers.0.bias', None),
                      ('model.diffusion_model.middle_block.0.in_layers.2.weight', (2, 3, 1, 0)),
                      ('model.diffusion_model.middle_block.0.in_layers.2.bias', None),
                      ('model.diffusion_model.middle_block.0.emb_layers.1.weight', (1, 0)),
                      ('model.diffusion_model.middle_block.0.emb_layers.1.bias', None),
                      ('model.diffusion_model.middle_block.0.out_layers.0.weight', None),
                      ('model.diffusion_model.middle_block.0.out_layers.0.bias', None),
                      ('model.diffusion_model.middle_block.0.out_layers.3.weight', (2, 3, 1, 0)),
                      ('model.diffusion_model.middle_block.0.out_layers.3.bias', None),
                      ('model.diffusion_model.middle_block.1.norm.weight', None),
                      ('model.diffusion_model.middle_block.1.norm.bias', None),
                      ('model.diffusion_model.middle_block.1.proj_in.weight', (2, 3, 1, 0)),
                      ('model.diffusion_model.middle_block.1.proj_in.bias', None),
                      ('model.diffusion_model.middle_block.1.transformer_blocks.0.norm1.weight', None),
                      ('model.diffusion_model.middle_block.1.transformer_blocks.0.norm1.bias', None),
                      ('model.diffusion_model.middle_block.1.transformer_blocks.0.attn1.to_q.weight', (1, 0)),
                      ('model.diffusion_model.middle_block.1.transformer_blocks.0.attn1.to_k.weight', (1, 0)),
                      ('model.diffusion_model.middle_block.1.transformer_blocks.0.attn1.to_v.weight', (1, 0)),
                      ('model.diffusion_model.middle_block.1.transformer_blocks.0.attn1.to_out.0.weight', (1, 0)),
                      ('model.diffusion_model.middle_block.1.transformer_blocks.0.attn1.to_out.0.bias', None),
                      ('model.diffusion_model.middle_block.1.transformer_blocks.0.norm2.weight', None),
                      ('model.diffusion_model.middle_block.1.transformer_blocks.0.norm2.bias', None),
                      ('model.diffusion_model.middle_block.1.transformer_blocks.0.attn2.to_q.weight', (1, 0)),
                      ('model.diffusion_model.middle_block.1.transformer_blocks.0.attn2.to_k.weight', (1, 0)),
                      ('model.diffusion_model.middle_block.1.transformer_blocks.0.attn2.to_v.weight', (1, 0)),
                      ('model.diffusion_model.middle_block.1.transformer_blocks.0.attn2.to_out.0.weight', (1, 0)),
                      ('model.diffusion_model.middle_block.1.transformer_blocks.0.attn2.to_out.0.bias', None),
                      ('model.diffusion_model.middle_block.1.transformer_blocks.0.norm3.weight', None),
                      ('model.diffusion_model.middle_block.1.transformer_blocks.0.norm3.bias', None),
                      ('model.diffusion_model.middle_block.1.transformer_blocks.0.ff.net.0.proj.weight', (1, 0)),
                      ('model.diffusion_model.middle_block.1.transformer_blocks.0.ff.net.0.proj.bias', None),
                      ('model.diffusion_model.middle_block.1.transformer_blocks.0.ff.net.2.weight', (1, 0)),
                      ('model.diffusion_model.middle_block.1.transformer_blocks.0.ff.net.2.bias', None),
                      ('model.diffusion_model.middle_block.1.proj_out.weight', (2, 3, 1, 0)),
                      ('model.diffusion_model.middle_block.1.proj_out.bias', None),
                      ('model.diffusion_model.middle_block.2.in_layers.0.weight', None),
                      ('model.diffusion_model.middle_block.2.in_layers.0.bias', None),
                      ('model.diffusion_model.middle_block.2.in_layers.2.weight', (2, 3, 1, 0)),
                      ('model.diffusion_model.middle_block.2.in_layers.2.bias', None),
                      ('model.diffusion_model.middle_block.2.emb_layers.1.weight', (1, 0)),
                      ('model.diffusion_model.middle_block.2.emb_layers.1.bias', None),
                      ('model.diffusion_model.middle_block.2.out_layers.0.weight', None),
                      ('model.diffusion_model.middle_block.2.out_layers.0.bias', None),
                      ('model.diffusion_model.middle_block.2.out_layers.3.weight', (2, 3, 1, 0)),
                      ('model.diffusion_model.middle_block.2.out_layers.3.bias', None),
                      ('model.diffusion_model.output_blocks.0.0.in_layers.0.weight', None),
                      ('model.diffusion_model.output_blocks.0.0.in_layers.0.bias', None),
                      ('model.diffusion_model.output_blocks.0.0.in_layers.2.weight', (2, 3, 1, 0)),
                      ('model.diffusion_model.output_blocks.0.0.in_layers.2.bias', None),
                      ('model.diffusion_model.output_blocks.0.0.emb_layers.1.weight', (1, 0)),
                      ('model.diffusion_model.output_blocks.0.0.emb_layers.1.bias', None),
                      ('model.diffusion_model.output_blocks.0.0.out_layers.0.weight', None),
                      ('model.diffusion_model.output_blocks.0.0.out_layers.0.bias', None),
                      ('model.diffusion_model.output_blocks.0.0.out_layers.3.weight', (2, 3, 1, 0)),
                      ('model.diffusion_model.output_blocks.0.0.out_layers.3.bias', None),
                      ('model.diffusion_model.output_blocks.0.0.skip_connection.weight', (2, 3, 1, 0)),
                      ('model.diffusion_model.output_blocks.0.0.skip_connection.bias', None),
                      ('model.diffusion_model.output_blocks.1.0.in_layers.0.weight', None),
                      ('model.diffusion_model.output_blocks.1.0.in_layers.0.bias', None),
                      ('model.diffusion_model.output_blocks.1.0.in_layers.2.weight', (2, 3, 1, 0)),
                      ('model.diffusion_model.output_blocks.1.0.in_layers.2.bias', None),
                      ('model.diffusion_model.output_blocks.1.0.emb_layers.1.weight', (1, 0)),
                      ('model.diffusion_model.output_blocks.1.0.emb_layers.1.bias', None),
                      ('model.diffusion_model.output_blocks.1.0.out_layers.0.weight', None),
                      ('model.diffusion_model.output_blocks.1.0.out_layers.0.bias', None),
                      ('model.diffusion_model.output_blocks.1.0.out_layers.3.weight', (2, 3, 1, 0)),
                      ('model.diffusion_model.output_blocks.1.0.out_layers.3.bias', None),
                      ('model.diffusion_model.output_blocks.1.0.skip_connection.weight', (2, 3, 1, 0)),
                      ('model.diffusion_model.output_blocks.1.0.skip_connection.bias', None),
                      ('model.diffusion_model.output_blocks.2.0.in_layers.0.weight', None),
                      ('model.diffusion_model.output_blocks.2.0.in_layers.0.bias', None),
                      ('model.diffusion_model.output_blocks.2.0.in_layers.2.weight', (2, 3, 1, 0)),
                      ('model.diffusion_model.output_blocks.2.0.in_layers.2.bias', None),
                      ('model.diffusion_model.output_blocks.2.0.emb_layers.1.weight', (1, 0)),
                      ('model.diffusion_model.output_blocks.2.0.emb_layers.1.bias', None),
                      ('model.diffusion_model.output_blocks.2.0.out_layers.0.weight', None),
                      ('model.diffusion_model.output_blocks.2.0.out_layers.0.bias', None),
                      ('model.diffusion_model.output_blocks.2.0.out_layers.3.weight', (2, 3, 1, 0)),
                      ('model.diffusion_model.output_blocks.2.0.out_layers.3.bias', None),
                      ('model.diffusion_model.output_blocks.2.0.skip_connection.weight', (2, 3, 1, 0)),
                      ('model.diffusion_model.output_blocks.2.0.skip_connection.bias', None),
                      ('model.diffusion_model.output_blocks.2.1.conv.weight', (2, 3, 1, 0)),
                      ('model.diffusion_model.output_blocks.2.1.conv.bias', None),
                      ('model.diffusion_model.output_blocks.3.0.in_layers.0.weight', None),
                      ('model.diffusion_model.output_blocks.3.0.in_layers.0.bias', None),
                      ('model.diffusion_model.output_blocks.3.0.in_layers.2.weight', (2, 3, 1, 0)),
                      ('model.diffusion_model.output_blocks.3.0.in_layers.2.bias', None),
                      ('model.diffusion_model.output_blocks.3.0.emb_layers.1.weight', (1, 0)),
                      ('model.diffusion_model.output_blocks.3.0.emb_layers.1.bias', None),
                      ('model.diffusion_model.output_blocks.3.0.out_layers.0.weight', None),
                      ('model.diffusion_model.output_blocks.3.0.out_layers.0.bias', None),
                      ('model.diffusion_model.output_blocks.3.0.out_layers.3.weight', (2, 3, 1, 0)),
                      ('model.diffusion_model.output_blocks.3.0.out_layers.3.bias', None),
                      ('model.diffusion_model.output_blocks.3.0.skip_connection.weight', (2, 3, 1, 0)),
                      ('model.diffusion_model.output_blocks.3.0.skip_connection.bias', None),
                      ('model.diffusion_model.output_blocks.3.1.norm.weight', None),
                      ('model.diffusion_model.output_blocks.3.1.norm.bias', None),
                      ('model.diffusion_model.output_blocks.3.1.proj_in.weight', (2, 3, 1, 0)),
                      ('model.diffusion_model.output_blocks.3.1.proj_in.bias', None),
                      ('model.diffusion_model.output_blocks.3.1.transformer_blocks.0.norm1.weight', None),
                      ('model.diffusion_model.output_blocks.3.1.transformer_blocks.0.norm1.bias', None),
                      ('model.diffusion_model.output_blocks.3.1.transformer_blocks.0.attn1.to_q.weight', (1, 0)),
                      ('model.diffusion_model.output_blocks.3.1.transformer_blocks.0.attn1.to_k.weight', (1, 0)),
                      ('model.diffusion_model.output_blocks.3.1.transformer_blocks.0.attn1.to_v.weight', (1, 0)),
                      ('model.diffusion_model.output_blocks.3.1.transformer_blocks.0.attn1.to_out.0.weight', (1, 0)),
                      ('model.diffusion_model.output_blocks.3.1.transformer_blocks.0.attn1.to_out.0.bias', None),
                      ('model.diffusion_model.output_blocks.3.1.transformer_blocks.0.norm2.weight', None),
                      ('model.diffusion_model.output_blocks.3.1.transformer_blocks.0.norm2.bias', None),
                      ('model.diffusion_model.output_blocks.3.1.transformer_blocks.0.attn2.to_q.weight', (1, 0)),
                      ('model.diffusion_model.output_blocks.3.1.transformer_blocks.0.attn2.to_k.weight', (1, 0)),
                      ('model.diffusion_model.output_blocks.3.1.transformer_blocks.0.attn2.to_v.weight', (1, 0)),
                      ('model.diffusion_model.output_blocks.3.1.transformer_blocks.0.attn2.to_out.0.weight', (1, 0)),
                      ('model.diffusion_model.output_blocks.3.1.transformer_blocks.0.attn2.to_out.0.bias', None),
                      ('model.diffusion_model.output_blocks.3.1.transformer_blocks.0.norm3.weight', None),
                      ('model.diffusion_model.output_blocks.3.1.transformer_blocks.0.norm3.bias', None),
                      ('model.diffusion_model.output_blocks.3.1.transformer_blocks.0.ff.net.0.proj.weight', (1, 0)),
                      ('model.diffusion_model.output_blocks.3.1.transformer_blocks.0.ff.net.0.proj.bias', None),
                      ('model.diffusion_model.output_blocks.3.1.transformer_blocks.0.ff.net.2.weight', (1, 0)),
                      ('model.diffusion_model.output_blocks.3.1.transformer_blocks.0.ff.net.2.bias', None),
                      ('model.diffusion_model.output_blocks.3.1.proj_out.weight', (2, 3, 1, 0)),
                      ('model.diffusion_model.output_blocks.3.1.proj_out.bias', None),
                      ('model.diffusion_model.output_blocks.4.0.in_layers.0.weight', None),
                      ('model.diffusion_model.output_blocks.4.0.in_layers.0.bias', None),
                      ('model.diffusion_model.output_blocks.4.0.in_layers.2.weight', (2, 3, 1, 0)),
                      ('model.diffusion_model.output_blocks.4.0.in_layers.2.bias', None),
                      ('model.diffusion_model.output_blocks.4.0.emb_layers.1.weight', (1, 0)),
                      ('model.diffusion_model.output_blocks.4.0.emb_layers.1.bias', None),
                      ('model.diffusion_model.output_blocks.4.0.out_layers.0.weight', None),
                      ('model.diffusion_model.output_blocks.4.0.out_layers.0.bias', None),
                      ('model.diffusion_model.output_blocks.4.0.out_layers.3.weight', (2, 3, 1, 0)),
                      ('model.diffusion_model.output_blocks.4.0.out_layers.3.bias', None),
                      ('model.diffusion_model.output_blocks.4.0.skip_connection.weight', (2, 3, 1, 0)),
                      ('model.diffusion_model.output_blocks.4.0.skip_connection.bias', None),
                      ('model.diffusion_model.output_blocks.4.1.norm.weight', None),
                      ('model.diffusion_model.output_blocks.4.1.norm.bias', None),
                      ('model.diffusion_model.output_blocks.4.1.proj_in.weight', (2, 3, 1, 0)),
                      ('model.diffusion_model.output_blocks.4.1.proj_in.bias', None),
                      ('model.diffusion_model.output_blocks.4.1.transformer_blocks.0.norm1.weight', None),
                      ('model.diffusion_model.output_blocks.4.1.transformer_blocks.0.norm1.bias', None),
                      ('model.diffusion_model.output_blocks.4.1.transformer_blocks.0.attn1.to_q.weight', (1, 0)),
                      ('model.diffusion_model.output_blocks.4.1.transformer_blocks.0.attn1.to_k.weight', (1, 0)),
                      ('model.diffusion_model.output_blocks.4.1.transformer_blocks.0.attn1.to_v.weight', (1, 0)),
                      ('model.diffusion_model.output_blocks.4.1.transformer_blocks.0.attn1.to_out.0.weight', (1, 0)),
                      ('model.diffusion_model.output_blocks.4.1.transformer_blocks.0.attn1.to_out.0.bias', None),
                      ('model.diffusion_model.output_blocks.4.1.transformer_blocks.0.norm2.weight', None),
                      ('model.diffusion_model.output_blocks.4.1.transformer_blocks.0.norm2.bias', None),
                      ('model.diffusion_model.output_blocks.4.1.transformer_blocks.0.attn2.to_q.weight', (1, 0)),
                      ('model.diffusion_model.output_blocks.4.1.transformer_blocks.0.attn2.to_k.weight', (1, 0)),
                      ('model.diffusion_model.output_blocks.4.1.transformer_blocks.0.attn2.to_v.weight', (1, 0)),
                      ('model.diffusion_model.output_blocks.4.1.transformer_blocks.0.attn2.to_out.0.weight', (1, 0)),
                      ('model.diffusion_model.output_blocks.4.1.transformer_blocks.0.attn2.to_out.0.bias', None),
                      ('model.diffusion_model.output_blocks.4.1.transformer_blocks.0.norm3.weight', None),
                      ('model.diffusion_model.output_blocks.4.1.transformer_blocks.0.norm3.bias', None),
                      ('model.diffusion_model.output_blocks.4.1.transformer_blocks.0.ff.net.0.proj.weight', (1, 0)),
                      ('model.diffusion_model.output_blocks.4.1.transformer_blocks.0.ff.net.0.proj.bias', None),
                      ('model.diffusion_model.output_blocks.4.1.transformer_blocks.0.ff.net.2.weight', (1, 0)),
                      ('model.diffusion_model.output_blocks.4.1.transformer_blocks.0.ff.net.2.bias', None),
                      ('model.diffusion_model.output_blocks.4.1.proj_out.weight', (2, 3, 1, 0)),
                      ('model.diffusion_model.output_blocks.4.1.proj_out.bias', None),
                      ('model.diffusion_model.output_blocks.5.0.in_layers.0.weight', None),
                      ('model.diffusion_model.output_blocks.5.0.in_layers.0.bias', None),
                      ('model.diffusion_model.output_blocks.5.0.in_layers.2.weight', (2, 3, 1, 0)),
                      ('model.diffusion_model.output_blocks.5.0.in_layers.2.bias', None),
                      ('model.diffusion_model.output_blocks.5.0.emb_layers.1.weight', (1, 0)),
                      ('model.diffusion_model.output_blocks.5.0.emb_layers.1.bias', None),
                      ('model.diffusion_model.output_blocks.5.0.out_layers.0.weight', None),
                      ('model.diffusion_model.output_blocks.5.0.out_layers.0.bias', None),
                      ('model.diffusion_model.output_blocks.5.0.out_layers.3.weight', (2, 3, 1, 0)),
                      ('model.diffusion_model.output_blocks.5.0.out_layers.3.bias', None),
                      ('model.diffusion_model.output_blocks.5.0.skip_connection.weight', (2, 3, 1, 0)),
                      ('model.diffusion_model.output_blocks.5.0.skip_connection.bias', None),
                      ('model.diffusion_model.output_blocks.5.1.norm.weight', None),
                      ('model.diffusion_model.output_blocks.5.1.norm.bias', None),
                      ('model.diffusion_model.output_blocks.5.1.proj_in.weight', (2, 3, 1, 0)),
                      ('model.diffusion_model.output_blocks.5.1.proj_in.bias', None),
                      ('model.diffusion_model.output_blocks.5.1.transformer_blocks.0.norm1.weight', None),
                      ('model.diffusion_model.output_blocks.5.1.transformer_blocks.0.norm1.bias', None),
                      ('model.diffusion_model.output_blocks.5.1.transformer_blocks.0.attn1.to_q.weight', (1, 0)),
                      ('model.diffusion_model.output_blocks.5.1.transformer_blocks.0.attn1.to_k.weight', (1, 0)),
                      ('model.diffusion_model.output_blocks.5.1.transformer_blocks.0.attn1.to_v.weight', (1, 0)),
                      ('model.diffusion_model.output_blocks.5.1.transformer_blocks.0.attn1.to_out.0.weight', (1, 0)),
                      ('model.diffusion_model.output_blocks.5.1.transformer_blocks.0.attn1.to_out.0.bias', None),
                      ('model.diffusion_model.output_blocks.5.1.transformer_blocks.0.norm2.weight', None),
                      ('model.diffusion_model.output_blocks.5.1.transformer_blocks.0.norm2.bias', None),
                      ('model.diffusion_model.output_blocks.5.1.transformer_blocks.0.attn2.to_q.weight', (1, 0)),
                      ('model.diffusion_model.output_blocks.5.1.transformer_blocks.0.attn2.to_k.weight', (1, 0)),
                      ('model.diffusion_model.output_blocks.5.1.transformer_blocks.0.attn2.to_v.weight', (1, 0)),
                      ('model.diffusion_model.output_blocks.5.1.transformer_blocks.0.attn2.to_out.0.weight', (1, 0)),
                      ('model.diffusion_model.output_blocks.5.1.transformer_blocks.0.attn2.to_out.0.bias', None),
                      ('model.diffusion_model.output_blocks.5.1.transformer_blocks.0.norm3.weight', None),
                      ('model.diffusion_model.output_blocks.5.1.transformer_blocks.0.norm3.bias', None),
                      ('model.diffusion_model.output_blocks.5.1.transformer_blocks.0.ff.net.0.proj.weight', (1, 0)),
                      ('model.diffusion_model.output_blocks.5.1.transformer_blocks.0.ff.net.0.proj.bias', None),
                      ('model.diffusion_model.output_blocks.5.1.transformer_blocks.0.ff.net.2.weight', (1, 0)),
                      ('model.diffusion_model.output_blocks.5.1.transformer_blocks.0.ff.net.2.bias', None),
                      ('model.diffusion_model.output_blocks.5.1.proj_out.weight', (2, 3, 1, 0)),
                      ('model.diffusion_model.output_blocks.5.1.proj_out.bias', None),
                      ('model.diffusion_model.output_blocks.5.2.conv.weight', (2, 3, 1, 0)),
                      ('model.diffusion_model.output_blocks.5.2.conv.bias', None),
                      ('model.diffusion_model.output_blocks.6.0.in_layers.0.weight', None),
                      ('model.diffusion_model.output_blocks.6.0.in_layers.0.bias', None),
                      ('model.diffusion_model.output_blocks.6.0.in_layers.2.weight', (2, 3, 1, 0)),
                      ('model.diffusion_model.output_blocks.6.0.in_layers.2.bias', None),
                      ('model.diffusion_model.output_blocks.6.0.emb_layers.1.weight', (1, 0)),
                      ('model.diffusion_model.output_blocks.6.0.emb_layers.1.bias', None),
                      ('model.diffusion_model.output_blocks.6.0.out_layers.0.weight', None),
                      ('model.diffusion_model.output_blocks.6.0.out_layers.0.bias', None),
                      ('model.diffusion_model.output_blocks.6.0.out_layers.3.weight', (2, 3, 1, 0)),
                      ('model.diffusion_model.output_blocks.6.0.out_layers.3.bias', None),
                      ('model.diffusion_model.output_blocks.6.0.skip_connection.weight', (2, 3, 1, 0)),
                      ('model.diffusion_model.output_blocks.6.0.skip_connection.bias', None),
                      ('model.diffusion_model.output_blocks.6.1.norm.weight', None),
                      ('model.diffusion_model.output_blocks.6.1.norm.bias', None),
                      ('model.diffusion_model.output_blocks.6.1.proj_in.weight', (2, 3, 1, 0)),
                      ('model.diffusion_model.output_blocks.6.1.proj_in.bias', None),
                      ('model.diffusion_model.output_blocks.6.1.transformer_blocks.0.norm1.weight', None),
                      ('model.diffusion_model.output_blocks.6.1.transformer_blocks.0.norm1.bias', None),
                      ('model.diffusion_model.output_blocks.6.1.transformer_blocks.0.attn1.to_q.weight', (1, 0)),
                      ('model.diffusion_model.output_blocks.6.1.transformer_blocks.0.attn1.to_k.weight', (1, 0)),
                      ('model.diffusion_model.output_blocks.6.1.transformer_blocks.0.attn1.to_v.weight', (1, 0)),
                      ('model.diffusion_model.output_blocks.6.1.transformer_blocks.0.attn1.to_out.0.weight', (1, 0)),
                      ('model.diffusion_model.output_blocks.6.1.transformer_blocks.0.attn1.to_out.0.bias', None),
                      ('model.diffusion_model.output_blocks.6.1.transformer_blocks.0.norm2.weight', None),
                      ('model.diffusion_model.output_blocks.6.1.transformer_blocks.0.norm2.bias', None),
                      ('model.diffusion_model.output_blocks.6.1.transformer_blocks.0.attn2.to_q.weight', (1, 0)),
                      ('model.diffusion_model.output_blocks.6.1.transformer_blocks.0.attn2.to_k.weight', (1, 0)),
                      ('model.diffusion_model.output_blocks.6.1.transformer_blocks.0.attn2.to_v.weight', (1, 0)),
                      ('model.diffusion_model.output_blocks.6.1.transformer_blocks.0.attn2.to_out.0.weight', (1, 0)),
                      ('model.diffusion_model.output_blocks.6.1.transformer_blocks.0.attn2.to_out.0.bias', None),
                      ('model.diffusion_model.output_blocks.6.1.transformer_blocks.0.norm3.weight', None),
                      ('model.diffusion_model.output_blocks.6.1.transformer_blocks.0.norm3.bias', None),
                      ('model.diffusion_model.output_blocks.6.1.transformer_blocks.0.ff.net.0.proj.weight', (1, 0)),
                      ('model.diffusion_model.output_blocks.6.1.transformer_blocks.0.ff.net.0.proj.bias', None),
                      ('model.diffusion_model.output_blocks.6.1.transformer_blocks.0.ff.net.2.weight', (1, 0)),
                      ('model.diffusion_model.output_blocks.6.1.transformer_blocks.0.ff.net.2.bias', None),
                      ('model.diffusion_model.output_blocks.6.1.proj_out.weight', (2, 3, 1, 0)),
                      ('model.diffusion_model.output_blocks.6.1.proj_out.bias', None),
                      ('model.diffusion_model.output_blocks.7.0.in_layers.0.weight', None),
                      ('model.diffusion_model.output_blocks.7.0.in_layers.0.bias', None),
                      ('model.diffusion_model.output_blocks.7.0.in_layers.2.weight', (2, 3, 1, 0)),
                      ('model.diffusion_model.output_blocks.7.0.in_layers.2.bias', None),
                      ('model.diffusion_model.output_blocks.7.0.emb_layers.1.weight', (1, 0)),
                      ('model.diffusion_model.output_blocks.7.0.emb_layers.1.bias', None),
                      ('model.diffusion_model.output_blocks.7.0.out_layers.0.weight', None),
                      ('model.diffusion_model.output_blocks.7.0.out_layers.0.bias', None),
                      ('model.diffusion_model.output_blocks.7.0.out_layers.3.weight', (2, 3, 1, 0)),
                      ('model.diffusion_model.output_blocks.7.0.out_layers.3.bias', None),
                      ('model.diffusion_model.output_blocks.7.0.skip_connection.weight', (2, 3, 1, 0)),
                      ('model.diffusion_model.output_blocks.7.0.skip_connection.bias', None),
                      ('model.diffusion_model.output_blocks.7.1.norm.weight', None),
                      ('model.diffusion_model.output_blocks.7.1.norm.bias', None),
                      ('model.diffusion_model.output_blocks.7.1.proj_in.weight', (2, 3, 1, 0)),
                      ('model.diffusion_model.output_blocks.7.1.proj_in.bias', None),
                      ('model.diffusion_model.output_blocks.7.1.transformer_blocks.0.norm1.weight', None),
                      ('model.diffusion_model.output_blocks.7.1.transformer_blocks.0.norm1.bias', None),
                      ('model.diffusion_model.output_blocks.7.1.transformer_blocks.0.attn1.to_q.weight', (1, 0)),
                      ('model.diffusion_model.output_blocks.7.1.transformer_blocks.0.attn1.to_k.weight', (1, 0)),
                      ('model.diffusion_model.output_blocks.7.1.transformer_blocks.0.attn1.to_v.weight', (1, 0)),
                      ('model.diffusion_model.output_blocks.7.1.transformer_blocks.0.attn1.to_out.0.weight', (1, 0)),
                      ('model.diffusion_model.output_blocks.7.1.transformer_blocks.0.attn1.to_out.0.bias', None),
                      ('model.diffusion_model.output_blocks.7.1.transformer_blocks.0.norm2.weight', None),
                      ('model.diffusion_model.output_blocks.7.1.transformer_blocks.0.norm2.bias', None),
                      ('model.diffusion_model.output_blocks.7.1.transformer_blocks.0.attn2.to_q.weight', (1, 0)),
                      ('model.diffusion_model.output_blocks.7.1.transformer_blocks.0.attn2.to_k.weight', (1, 0)),
                      ('model.diffusion_model.output_blocks.7.1.transformer_blocks.0.attn2.to_v.weight', (1, 0)),
                      ('model.diffusion_model.output_blocks.7.1.transformer_blocks.0.attn2.to_out.0.weight', (1, 0)),
                      ('model.diffusion_model.output_blocks.7.1.transformer_blocks.0.attn2.to_out.0.bias', None),
                      ('model.diffusion_model.output_blocks.7.1.transformer_blocks.0.norm3.weight', None),
                      ('model.diffusion_model.output_blocks.7.1.transformer_blocks.0.norm3.bias', None),
                      ('model.diffusion_model.output_blocks.7.1.transformer_blocks.0.ff.net.0.proj.weight', (1, 0)),
                      ('model.diffusion_model.output_blocks.7.1.transformer_blocks.0.ff.net.0.proj.bias', None),
                      ('model.diffusion_model.output_blocks.7.1.transformer_blocks.0.ff.net.2.weight', (1, 0)),
                      ('model.diffusion_model.output_blocks.7.1.transformer_blocks.0.ff.net.2.bias', None),
                      ('model.diffusion_model.output_blocks.7.1.proj_out.weight', (2, 3, 1, 0)),
                      ('model.diffusion_model.output_blocks.7.1.proj_out.bias', None),
                      ('model.diffusion_model.output_blocks.8.0.in_layers.0.weight', None),
                      ('model.diffusion_model.output_blocks.8.0.in_layers.0.bias', None),
                      ('model.diffusion_model.output_blocks.8.0.in_layers.2.weight', (2, 3, 1, 0)),
                      ('model.diffusion_model.output_blocks.8.0.in_layers.2.bias', None),
                      ('model.diffusion_model.output_blocks.8.0.emb_layers.1.weight', (1, 0)),
                      ('model.diffusion_model.output_blocks.8.0.emb_layers.1.bias', None),
                      ('model.diffusion_model.output_blocks.8.0.out_layers.0.weight', None),
                      ('model.diffusion_model.output_blocks.8.0.out_layers.0.bias', None),
                      ('model.diffusion_model.output_blocks.8.0.out_layers.3.weight', (2, 3, 1, 0)),
                      ('model.diffusion_model.output_blocks.8.0.out_layers.3.bias', None),
                      ('model.diffusion_model.output_blocks.8.0.skip_connection.weight', (2, 3, 1, 0)),
                      ('model.diffusion_model.output_blocks.8.0.skip_connection.bias', None),
                      ('model.diffusion_model.output_blocks.8.1.norm.weight', None),
                      ('model.diffusion_model.output_blocks.8.1.norm.bias', None),
                      ('model.diffusion_model.output_blocks.8.1.proj_in.weight', (2, 3, 1, 0)),
                      ('model.diffusion_model.output_blocks.8.1.proj_in.bias', None),
                      ('model.diffusion_model.output_blocks.8.1.transformer_blocks.0.norm1.weight', None),
                      ('model.diffusion_model.output_blocks.8.1.transformer_blocks.0.norm1.bias', None),
                      ('model.diffusion_model.output_blocks.8.1.transformer_blocks.0.attn1.to_q.weight', (1, 0)),
                      ('model.diffusion_model.output_blocks.8.1.transformer_blocks.0.attn1.to_k.weight', (1, 0)),
                      ('model.diffusion_model.output_blocks.8.1.transformer_blocks.0.attn1.to_v.weight', (1, 0)),
                      ('model.diffusion_model.output_blocks.8.1.transformer_blocks.0.attn1.to_out.0.weight', (1, 0)),
                      ('model.diffusion_model.output_blocks.8.1.transformer_blocks.0.attn1.to_out.0.bias', None),
                      ('model.diffusion_model.output_blocks.8.1.transformer_blocks.0.norm2.weight', None),
                      ('model.diffusion_model.output_blocks.8.1.transformer_blocks.0.norm2.bias', None),
                      ('model.diffusion_model.output_blocks.8.1.transformer_blocks.0.attn2.to_q.weight', (1, 0)),
                      ('model.diffusion_model.output_blocks.8.1.transformer_blocks.0.attn2.to_k.weight', (1, 0)),
                      ('model.diffusion_model.output_blocks.8.1.transformer_blocks.0.attn2.to_v.weight', (1, 0)),
                      ('model.diffusion_model.output_blocks.8.1.transformer_blocks.0.attn2.to_out.0.weight', (1, 0)),
                      ('model.diffusion_model.output_blocks.8.1.transformer_blocks.0.attn2.to_out.0.bias', None),
                      ('model.diffusion_model.output_blocks.8.1.transformer_blocks.0.norm3.weight', None),
                      ('model.diffusion_model.output_blocks.8.1.transformer_blocks.0.norm3.bias', None),
                      ('model.diffusion_model.output_blocks.8.1.transformer_blocks.0.ff.net.0.proj.weight', (1, 0)),
                      ('model.diffusion_model.output_blocks.8.1.transformer_blocks.0.ff.net.0.proj.bias', None),
                      ('model.diffusion_model.output_blocks.8.1.transformer_blocks.0.ff.net.2.weight', (1, 0)),
                      ('model.diffusion_model.output_blocks.8.1.transformer_blocks.0.ff.net.2.bias', None),
                      ('model.diffusion_model.output_blocks.8.1.proj_out.weight', (2, 3, 1, 0)),
                      ('model.diffusion_model.output_blocks.8.1.proj_out.bias', None),
                      ('model.diffusion_model.output_blocks.8.2.conv.weight', (2, 3, 1, 0)),
                      ('model.diffusion_model.output_blocks.8.2.conv.bias', None),
                      ('model.diffusion_model.output_blocks.9.0.in_layers.0.weight', None),
                      ('model.diffusion_model.output_blocks.9.0.in_layers.0.bias', None),
                      ('model.diffusion_model.output_blocks.9.0.in_layers.2.weight', (2, 3, 1, 0)),
                      ('model.diffusion_model.output_blocks.9.0.in_layers.2.bias', None),
                      ('model.diffusion_model.output_blocks.9.0.emb_layers.1.weight', (1, 0)),
                      ('model.diffusion_model.output_blocks.9.0.emb_layers.1.bias', None),
                      ('model.diffusion_model.output_blocks.9.0.out_layers.0.weight', None),
                      ('model.diffusion_model.output_blocks.9.0.out_layers.0.bias', None),
                      ('model.diffusion_model.output_blocks.9.0.out_layers.3.weight', (2, 3, 1, 0)),
                      ('model.diffusion_model.output_blocks.9.0.out_layers.3.bias', None),
                      ('model.diffusion_model.output_blocks.9.0.skip_connection.weight', (2, 3, 1, 0)),
                      ('model.diffusion_model.output_blocks.9.0.skip_connection.bias', None),
                      ('model.diffusion_model.output_blocks.9.1.norm.weight', None),
                      ('model.diffusion_model.output_blocks.9.1.norm.bias', None),
                      ('model.diffusion_model.output_blocks.9.1.proj_in.weight', (2, 3, 1, 0)),
                      ('model.diffusion_model.output_blocks.9.1.proj_in.bias', None),
                      ('model.diffusion_model.output_blocks.9.1.transformer_blocks.0.norm1.weight', None),
                      ('model.diffusion_model.output_blocks.9.1.transformer_blocks.0.norm1.bias', None),
                      ('model.diffusion_model.output_blocks.9.1.transformer_blocks.0.attn1.to_q.weight', (1, 0)),
                      ('model.diffusion_model.output_blocks.9.1.transformer_blocks.0.attn1.to_k.weight', (1, 0)),
                      ('model.diffusion_model.output_blocks.9.1.transformer_blocks.0.attn1.to_v.weight', (1, 0)),
                      ('model.diffusion_model.output_blocks.9.1.transformer_blocks.0.attn1.to_out.0.weight', (1, 0)),
                      ('model.diffusion_model.output_blocks.9.1.transformer_blocks.0.attn1.to_out.0.bias', None),
                      ('model.diffusion_model.output_blocks.9.1.transformer_blocks.0.norm2.weight', None),
                      ('model.diffusion_model.output_blocks.9.1.transformer_blocks.0.norm2.bias', None),
                      ('model.diffusion_model.output_blocks.9.1.transformer_blocks.0.attn2.to_q.weight', (1, 0)),
                      ('model.diffusion_model.output_blocks.9.1.transformer_blocks.0.attn2.to_k.weight', (1, 0)),
                      ('model.diffusion_model.output_blocks.9.1.transformer_blocks.0.attn2.to_v.weight', (1, 0)),
                      ('model.diffusion_model.output_blocks.9.1.transformer_blocks.0.attn2.to_out.0.weight', (1, 0)),
                      ('model.diffusion_model.output_blocks.9.1.transformer_blocks.0.attn2.to_out.0.bias', None),
                      ('model.diffusion_model.output_blocks.9.1.transformer_blocks.0.norm3.weight', None),
                      ('model.diffusion_model.output_blocks.9.1.transformer_blocks.0.norm3.bias', None),
                      ('model.diffusion_model.output_blocks.9.1.transformer_blocks.0.ff.net.0.proj.weight', (1, 0)),
                      ('model.diffusion_model.output_blocks.9.1.transformer_blocks.0.ff.net.0.proj.bias', None),
                      ('model.diffusion_model.output_blocks.9.1.transformer_blocks.0.ff.net.2.weight', (1, 0)),
                      ('model.diffusion_model.output_blocks.9.1.transformer_blocks.0.ff.net.2.bias', None),
                      ('model.diffusion_model.output_blocks.9.1.proj_out.weight', (2, 3, 1, 0)),
                      ('model.diffusion_model.output_blocks.9.1.proj_out.bias', None),
                      ('model.diffusion_model.output_blocks.10.0.in_layers.0.weight', None),
                      ('model.diffusion_model.output_blocks.10.0.in_layers.0.bias', None),
                      ('model.diffusion_model.output_blocks.10.0.in_layers.2.weight', (2, 3, 1, 0)),
                      ('model.diffusion_model.output_blocks.10.0.in_layers.2.bias', None),
                      ('model.diffusion_model.output_blocks.10.0.emb_layers.1.weight', (1, 0)),
                      ('model.diffusion_model.output_blocks.10.0.emb_layers.1.bias', None),
                      ('model.diffusion_model.output_blocks.10.0.out_layers.0.weight', None),
                      ('model.diffusion_model.output_blocks.10.0.out_layers.0.bias', None),
                      ('model.diffusion_model.output_blocks.10.0.out_layers.3.weight', (2, 3, 1, 0)),
                      ('model.diffusion_model.output_blocks.10.0.out_layers.3.bias', None),
                      ('model.diffusion_model.output_blocks.10.0.skip_connection.weight', (2, 3, 1, 0)),
                      ('model.diffusion_model.output_blocks.10.0.skip_connection.bias', None),
                      ('model.diffusion_model.output_blocks.10.1.norm.weight', None),
                      ('model.diffusion_model.output_blocks.10.1.norm.bias', None),
                      ('model.diffusion_model.output_blocks.10.1.proj_in.weight', (2, 3, 1, 0)),
                      ('model.diffusion_model.output_blocks.10.1.proj_in.bias', None),
                      ('model.diffusion_model.output_blocks.10.1.transformer_blocks.0.norm1.weight', None),
                      ('model.diffusion_model.output_blocks.10.1.transformer_blocks.0.norm1.bias', None),
                      ('model.diffusion_model.output_blocks.10.1.transformer_blocks.0.attn1.to_q.weight', (1, 0)),
                      ('model.diffusion_model.output_blocks.10.1.transformer_blocks.0.attn1.to_k.weight', (1, 0)),
                      ('model.diffusion_model.output_blocks.10.1.transformer_blocks.0.attn1.to_v.weight', (1, 0)),
                      ('model.diffusion_model.output_blocks.10.1.transformer_blocks.0.attn1.to_out.0.weight', (1, 0)),
                      ('model.diffusion_model.output_blocks.10.1.transformer_blocks.0.attn1.to_out.0.bias', None),
                      ('model.diffusion_model.output_blocks.10.1.transformer_blocks.0.norm2.weight', None),
                      ('model.diffusion_model.output_blocks.10.1.transformer_blocks.0.norm2.bias', None),
                      ('model.diffusion_model.output_blocks.10.1.transformer_blocks.0.attn2.to_q.weight', (1, 0)),
                      ('model.diffusion_model.output_blocks.10.1.transformer_blocks.0.attn2.to_k.weight', (1, 0)),
                      ('model.diffusion_model.output_blocks.10.1.transformer_blocks.0.attn2.to_v.weight', (1, 0)),
                      ('model.diffusion_model.output_blocks.10.1.transformer_blocks.0.attn2.to_out.0.weight', (1, 0)),
                      ('model.diffusion_model.output_blocks.10.1.transformer_blocks.0.attn2.to_out.0.bias', None),
                      ('model.diffusion_model.output_blocks.10.1.transformer_blocks.0.norm3.weight', None),
                      ('model.diffusion_model.output_blocks.10.1.transformer_blocks.0.norm3.bias', None),
                      ('model.diffusion_model.output_blocks.10.1.transformer_blocks.0.ff.net.0.proj.weight', (1, 0)),
                      ('model.diffusion_model.output_blocks.10.1.transformer_blocks.0.ff.net.0.proj.bias', None),
                      ('model.diffusion_model.output_blocks.10.1.transformer_blocks.0.ff.net.2.weight', (1, 0)),
                      ('model.diffusion_model.output_blocks.10.1.transformer_blocks.0.ff.net.2.bias', None),
                      ('model.diffusion_model.output_blocks.10.1.proj_out.weight', (2, 3, 1, 0)),
                      ('model.diffusion_model.output_blocks.10.1.proj_out.bias', None),
                      ('model.diffusion_model.output_blocks.11.0.in_layers.0.weight', None),
                      ('model.diffusion_model.output_blocks.11.0.in_layers.0.bias', None),
                      ('model.diffusion_model.output_blocks.11.0.in_layers.2.weight', (2, 3, 1, 0)),
                      ('model.diffusion_model.output_blocks.11.0.in_layers.2.bias', None),
                      ('model.diffusion_model.output_blocks.11.0.emb_layers.1.weight', (1, 0)),
                      ('model.diffusion_model.output_blocks.11.0.emb_layers.1.bias', None),
                      ('model.diffusion_model.output_blocks.11.0.out_layers.0.weight', None),
                      ('model.diffusion_model.output_blocks.11.0.out_layers.0.bias', None),
                      ('model.diffusion_model.output_blocks.11.0.out_layers.3.weight', (2, 3, 1, 0)),
                      ('model.diffusion_model.output_blocks.11.0.out_layers.3.bias', None),
                      ('model.diffusion_model.output_blocks.11.0.skip_connection.weight', (2, 3, 1, 0)),
                      ('model.diffusion_model.output_blocks.11.0.skip_connection.bias', None),
                      ('model.diffusion_model.output_blocks.11.1.norm.weight', None),
                      ('model.diffusion_model.output_blocks.11.1.norm.bias', None),
                      ('model.diffusion_model.output_blocks.11.1.proj_in.weight', (2, 3, 1, 0)),
                      ('model.diffusion_model.output_blocks.11.1.proj_in.bias', None),
                      ('model.diffusion_model.output_blocks.11.1.transformer_blocks.0.norm1.weight', None),
                      ('model.diffusion_model.output_blocks.11.1.transformer_blocks.0.norm1.bias', None),
                      ('model.diffusion_model.output_blocks.11.1.transformer_blocks.0.attn1.to_q.weight', (1, 0)),
                      ('model.diffusion_model.output_blocks.11.1.transformer_blocks.0.attn1.to_k.weight', (1, 0)),
                      ('model.diffusion_model.output_blocks.11.1.transformer_blocks.0.attn1.to_v.weight', (1, 0)),
                      ('model.diffusion_model.output_blocks.11.1.transformer_blocks.0.attn1.to_out.0.weight', (1, 0)),
                      ('model.diffusion_model.output_blocks.11.1.transformer_blocks.0.attn1.to_out.0.bias', None),
                      ('model.diffusion_model.output_blocks.11.1.transformer_blocks.0.norm2.weight', None),
                      ('model.diffusion_model.output_blocks.11.1.transformer_blocks.0.norm2.bias', None),
                      ('model.diffusion_model.output_blocks.11.1.transformer_blocks.0.attn2.to_q.weight', (1, 0)),
                      ('model.diffusion_model.output_blocks.11.1.transformer_blocks.0.attn2.to_k.weight', (1, 0)),
                      ('model.diffusion_model.output_blocks.11.1.transformer_blocks.0.attn2.to_v.weight', (1, 0)),
                      ('model.diffusion_model.output_blocks.11.1.transformer_blocks.0.attn2.to_out.0.weight', (1, 0)),
                      ('model.diffusion_model.output_blocks.11.1.transformer_blocks.0.attn2.to_out.0.bias', None),
                      ('model.diffusion_model.output_blocks.11.1.transformer_blocks.0.norm3.weight', None),
                      ('model.diffusion_model.output_blocks.11.1.transformer_blocks.0.norm3.bias', None),
                      ('model.diffusion_model.output_blocks.11.1.transformer_blocks.0.ff.net.0.proj.weight', (1, 0)),
                      ('model.diffusion_model.output_blocks.11.1.transformer_blocks.0.ff.net.0.proj.bias', None),
                      ('model.diffusion_model.output_blocks.11.1.transformer_blocks.0.ff.net.2.weight', (1, 0)),
                      ('model.diffusion_model.output_blocks.11.1.transformer_blocks.0.ff.net.2.bias', None),
                      ('model.diffusion_model.output_blocks.11.1.proj_out.weight', (2, 3, 1, 0)),
                      ('model.diffusion_model.output_blocks.11.1.proj_out.bias', None),
                      ('model.diffusion_model.out.0.weight', None),
                      ('model.diffusion_model.out.0.bias', None),
                      ('model.diffusion_model.out.2.weight', (2, 3, 1, 0)),
                      ('model.diffusion_model.out.2.bias', None)],
    'encoder': [
        ('encoder.conv_in.weight', (2, 3, 1, 0)),
        ('encoder.conv_in.bias', None),
        ('encoder.down_blocks.0.resnets.0.norm1.weight', None),
        ('encoder.down_blocks.0.resnets.0.norm1.bias', None),
        ('encoder.down_blocks.0.resnets.0.conv1.weight', (2, 3, 1, 0)),
        ('encoder.down_blocks.0.resnets.0.conv1.bias', None),
        ('encoder.down_blocks.0.resnets.0.norm2.weight', None),
        ('encoder.down_blocks.0.resnets.0.norm2.bias', None),
        ('encoder.down_blocks.0.resnets.0.conv2.weight', (2, 3, 1, 0)),
        ('encoder.down_blocks.0.resnets.0.conv2.bias', None),
        ('encoder.down_blocks.0.resnets.1.norm1.weight', None),
        ('encoder.down_blocks.0.resnets.1.norm1.bias', None),
        ('encoder.down_blocks.0.resnets.1.conv1.weight', (2, 3, 1, 0)),
        ('encoder.down_blocks.0.resnets.1.conv1.bias', None),
        ('encoder.down_blocks.0.resnets.1.norm2.weight', None),
        ('encoder.down_blocks.0.resnets.1.norm2.bias', None),
        ('encoder.down_blocks.0.resnets.1.conv2.weight', (2, 3, 1, 0)),
        ('encoder.down_blocks.0.resnets.1.conv2.bias', None),
        ('encoder.down_blocks.0.downsamplers.0.conv.weight', (2, 3, 1, 0)),
        ('encoder.down_blocks.0.downsamplers.0.conv.bias', None),
        ('encoder.down_blocks.1.resnets.0.norm1.weight', None),
        ('encoder.down_blocks.1.resnets.0.norm1.bias', None),
        ('encoder.down_blocks.1.resnets.0.conv1.weight', (2, 3, 1, 0)),
        ('encoder.down_blocks.1.resnets.0.conv1.bias', None),
        ('encoder.down_blocks.1.resnets.0.norm2.weight', None),
        ('encoder.down_blocks.1.resnets.0.norm2.bias', None),
        ('encoder.down_blocks.1.resnets.0.conv2.weight', (2, 3, 1, 0)),
        ('encoder.down_blocks.1.resnets.0.conv2.bias', None),
        ('encoder.down_blocks.1.resnets.0.conv_shortcut.weight', (2, 3, 1, 0)),
        ('encoder.down_blocks.1.resnets.0.conv_shortcut.bias', None),
        ('encoder.down_blocks.1.resnets.1.norm1.weight', None),
        ('encoder.down_blocks.1.resnets.1.norm1.bias', None),
        ('encoder.down_blocks.1.resnets.1.conv1.weight', (2, 3, 1, 0)),
        ('encoder.down_blocks.1.resnets.1.conv1.bias', None),
        ('encoder.down_blocks.1.resnets.1.norm2.weight', None),
        ('encoder.down_blocks.1.resnets.1.norm2.bias', None),
        ('encoder.down_blocks.1.resnets.1.conv2.weight', (2, 3, 1, 0)),
        ('encoder.down_blocks.1.resnets.1.conv2.bias', None),
        ('encoder.down_blocks.1.downsamplers.0.conv.weight', (2, 3, 1, 0)),
        ('encoder.down_blocks.1.downsamplers.0.conv.bias', None),
        ('encoder.down_blocks.2.resnets.0.norm1.weight', None),
        ('encoder.down_blocks.2.resnets.0.norm1.bias', None),
        ('encoder.down_blocks.2.resnets.0.conv1.weight', (2, 3, 1, 0)),
        ('encoder.down_blocks.2.resnets.0.conv1.bias', None),
        ('encoder.down_blocks.2.resnets.0.norm2.weight', None),
        ('encoder.down_blocks.2.resnets.0.norm2.bias', None),
        ('encoder.down_blocks.2.resnets.0.conv2.weight', (2, 3, 1, 0)),
        ('encoder.down_blocks.2.resnets.0.conv2.bias', None),
        ('encoder.down_blocks.2.resnets.0.conv_shortcut.weight', (2, 3, 1, 0)),
        ('encoder.down_blocks.2.resnets.0.conv_shortcut.bias', None),
        ('encoder.down_blocks.2.resnets.1.norm1.weight', None),
        ('encoder.down_blocks.2.resnets.1.norm1.bias', None),
        ('encoder.down_blocks.2.resnets.1.conv1.weight', (2, 3, 1, 0)),
        ('encoder.down_blocks.2.resnets.1.conv1.bias', None),
        ('encoder.down_blocks.2.resnets.1.norm2.weight', None),
        ('encoder.down_blocks.2.resnets.1.norm2.bias', None),
        ('encoder.down_blocks.2.resnets.1.conv2.weight', (2, 3, 1, 0)),
        ('encoder.down_blocks.2.resnets.1.conv2.bias', None),
        ('encoder.down_blocks.2.downsamplers.0.conv.weight', (2, 3, 1, 0)),
        ('encoder.down_blocks.2.downsamplers.0.conv.bias', None),
        ('encoder.down_blocks.3.resnets.0.norm1.weight', None),
        ('encoder.down_blocks.3.resnets.0.norm1.bias', None),
        ('encoder.down_blocks.3.resnets.0.conv1.weight', (2, 3, 1, 0)),
        ('encoder.down_blocks.3.resnets.0.conv1.bias', None),
        ('encoder.down_blocks.3.resnets.0.norm2.weight', None),
        ('encoder.down_blocks.3.resnets.0.norm2.bias', None),
        ('encoder.down_blocks.3.resnets.0.conv2.weight', (2, 3, 1, 0)),
        ('encoder.down_blocks.3.resnets.0.conv2.bias', None),
        ('encoder.down_blocks.3.resnets.1.norm1.weight', None),
        ('encoder.down_blocks.3.resnets.1.norm1.bias', None),
        ('encoder.down_blocks.3.resnets.1.conv1.weight', (2, 3, 1, 0)),
        ('encoder.down_blocks.3.resnets.1.conv1.bias', None),
        ('encoder.down_blocks.3.resnets.1.norm2.weight', None),
        ('encoder.down_blocks.3.resnets.1.norm2.bias', None),
        ('encoder.down_blocks.3.resnets.1.conv2.weight', (2, 3, 1, 0)),
        ('encoder.down_blocks.3.resnets.1.conv2.bias', None),
        ('encoder.mid_block.resnets.0.norm1.weight', None),
        ('encoder.mid_block.resnets.0.norm1.bias', None),
        ('encoder.mid_block.resnets.0.conv1.weight', (2, 3, 1, 0)),
        ('encoder.mid_block.resnets.0.conv1.bias', None),
        ('encoder.mid_block.resnets.0.norm2.weight', None),
        ('encoder.mid_block.resnets.0.norm2.bias', None),
        ('encoder.mid_block.resnets.0.conv2.weight', (2, 3, 1, 0)),
        ('encoder.mid_block.resnets.0.conv2.bias', None),
        ('encoder.mid_block.attentions.0.group_norm.weight', None),
        ('encoder.mid_block.attentions.0.group_norm.bias', None),
        ('encoder.mid_block.attentions.0.query.weight', (1, 0)),
        ('encoder.mid_block.attentions.0.query.bias', None),
        ('encoder.mid_block.attentions.0.key.weight', (1, 0)),
        ('encoder.mid_block.attentions.0.key.bias', None),
        ('encoder.mid_block.attentions.0.value.weight', (1, 0)),
        ('encoder.mid_block.attentions.0.value.bias', None),
        ('encoder.mid_block.attentions.0.proj_attn.weight', (1, 0)),
        ('encoder.mid_block.attentions.0.proj_attn.bias', None),
        ('encoder.mid_block.resnets.1.norm1.weight', None),
        ('encoder.mid_block.resnets.1.norm1.bias', None),
        ('encoder.mid_block.resnets.1.conv1.weight', (2, 3, 1, 0)),
        ('encoder.mid_block.resnets.1.conv1.bias', None),
        ('encoder.mid_block.resnets.1.norm2.weight', None),
        ('encoder.mid_block.resnets.1.norm2.bias', None),
        ('encoder.mid_block.resnets.1.conv2.weight', (2, 3, 1, 0)),
        ('encoder.mid_block.resnets.1.conv2.bias', None),
        ('encoder.conv_norm_out.weight', None),
        ('encoder.conv_norm_out.bias', None),
        ('encoder.conv_out.weight', (2, 3, 1, 0)),
        ('encoder.conv_out.bias', None),
        ('quant_conv.weight', (2, 3, 1, 0)),
        ('quant_conv.bias', None),
    ],
    'decoder': [
        ('post_quant_conv.weight', (2, 3, 1, 0)),
        ('post_quant_conv.bias', None),
        ('decoder.conv_in.weight', (2, 3, 1, 0)),
        ('decoder.conv_in.bias', None),
        ('decoder.mid_block.resnets.0.norm1.weight', None),
        ('decoder.mid_block.resnets.0.norm1.bias', None),
        ('decoder.mid_block.resnets.0.conv1.weight', (2, 3, 1, 0)),
        ('decoder.mid_block.resnets.0.conv1.bias', None),
        ('decoder.mid_block.resnets.0.norm2.weight', None),
        ('decoder.mid_block.resnets.0.norm2.bias', None),
        ('decoder.mid_block.resnets.0.conv2.weight', (2, 3, 1, 0)),
        ('decoder.mid_block.resnets.0.conv2.bias', None),
        ('decoder.mid_block.attentions.0.group_norm.weight', None),
        ('decoder.mid_block.attentions.0.group_norm.bias', None),
        ('decoder.mid_block.attentions.0.query.weight', (1, 0)),
        ('decoder.mid_block.attentions.0.query.bias', None),
        ('decoder.mid_block.attentions.0.key.weight', (1, 0)),
        ('decoder.mid_block.attentions.0.key.bias', None),
        ('decoder.mid_block.attentions.0.value.weight', (1, 0)),
        ('decoder.mid_block.attentions.0.value.bias', None),
        ('decoder.mid_block.attentions.0.proj_attn.weight', (1, 0)),
        ('decoder.mid_block.attentions.0.proj_attn.bias', None),
        ('decoder.mid_block.resnets.1.norm1.weight', None),
        ('decoder.mid_block.resnets.1.norm1.bias', None),
        ('decoder.mid_block.resnets.1.conv1.weight', (2, 3, 1, 0)),
        ('decoder.mid_block.resnets.1.conv1.bias', None),
        ('decoder.mid_block.resnets.1.norm2.weight', None),
        ('decoder.mid_block.resnets.1.norm2.bias', None),
        ('decoder.mid_block.resnets.1.conv2.weight', (2, 3, 1, 0)),
        ('decoder.mid_block.resnets.1.conv2.bias', None),
        ('decoder.up_blocks.0.resnets.0.norm1.weight', None),
        ('decoder.up_blocks.0.resnets.0.norm1.bias', None),
        ('decoder.up_blocks.0.resnets.0.conv1.weight', (2, 3, 1, 0)),
        ('decoder.up_blocks.0.resnets.0.conv1.bias', None),
        ('decoder.up_blocks.0.resnets.0.norm2.weight', None),
        ('decoder.up_blocks.0.resnets.0.norm2.bias', None),
        ('decoder.up_blocks.0.resnets.0.conv2.weight', (2, 3, 1, 0)),
        ('decoder.up_blocks.0.resnets.0.conv2.bias', None),
        ('decoder.up_blocks.0.resnets.1.norm1.weight', None),
        ('decoder.up_blocks.0.resnets.1.norm1.bias', None),
        ('decoder.up_blocks.0.resnets.1.conv1.weight', (2, 3, 1, 0)),
        ('decoder.up_blocks.0.resnets.1.conv1.bias', None),
        ('decoder.up_blocks.0.resnets.1.norm2.weight', None),
        ('decoder.up_blocks.0.resnets.1.norm2.bias', None),
        ('decoder.up_blocks.0.resnets.1.conv2.weight', (2, 3, 1, 0)),
        ('decoder.up_blocks.0.resnets.1.conv2.bias', None),
        ('decoder.up_blocks.0.resnets.2.norm1.weight', None),
        ('decoder.up_blocks.0.resnets.2.norm1.bias', None),
        ('decoder.up_blocks.0.resnets.2.conv1.weight', (2, 3, 1, 0)),
        ('decoder.up_blocks.0.resnets.2.conv1.bias', None),
        ('decoder.up_blocks.0.resnets.2.norm2.weight', None),
        ('decoder.up_blocks.0.resnets.2.norm2.bias', None),
        ('decoder.up_blocks.0.resnets.2.conv2.weight', (2, 3, 1, 0)),
        ('decoder.up_blocks.0.resnets.2.conv2.bias', None),
        ('decoder.up_blocks.0.upsamplers.0.conv.weight', (2, 3, 1, 0)),
        ('decoder.up_blocks.0.upsamplers.0.conv.bias', None),
        ('decoder.up_blocks.1.resnets.0.norm1.weight', None),
        ('decoder.up_blocks.1.resnets.0.norm1.bias', None),
        ('decoder.up_blocks.1.resnets.0.conv1.weight', (2, 3, 1, 0)),
        ('decoder.up_blocks.1.resnets.0.conv1.bias', None),
        ('decoder.up_blocks.1.resnets.0.norm2.weight', None),
        ('decoder.up_blocks.1.resnets.0.norm2.bias', None),
        ('decoder.up_blocks.1.resnets.0.conv2.weight', (2, 3, 1, 0)),
        ('decoder.up_blocks.1.resnets.0.conv2.bias', None),
        ('decoder.up_blocks.1.resnets.1.norm1.weight', None),
        ('decoder.up_blocks.1.resnets.1.norm1.bias', None),
        ('decoder.up_blocks.1.resnets.1.conv1.weight', (2, 3, 1, 0)),
        ('decoder.up_blocks.1.resnets.1.conv1.bias', None),
        ('decoder.up_blocks.1.resnets.1.norm2.weight', None),
        ('decoder.up_blocks.1.resnets.1.norm2.bias', None),
        ('decoder.up_blocks.1.resnets.1.conv2.weight', (2, 3, 1, 0)),
        ('decoder.up_blocks.1.resnets.1.conv2.bias', None),
        ('decoder.up_blocks.1.resnets.2.norm1.weight', None),
        ('decoder.up_blocks.1.resnets.2.norm1.bias', None),
        ('decoder.up_blocks.1.resnets.2.conv1.weight', (2, 3, 1, 0)),
        ('decoder.up_blocks.1.resnets.2.conv1.bias', None),
        ('decoder.up_blocks.1.resnets.2.norm2.weight', None),
        ('decoder.up_blocks.1.resnets.2.norm2.bias', None),
        ('decoder.up_blocks.1.resnets.2.conv2.weight', (2, 3, 1, 0)),
        ('decoder.up_blocks.1.resnets.2.conv2.bias', None),
        ('decoder.up_blocks.1.upsamplers.0.conv.weight', (2, 3, 1, 0)),
        ('decoder.up_blocks.1.upsamplers.0.conv.bias', None),
        ('decoder.up_blocks.2.resnets.0.norm1.weight', None),
        ('decoder.up_blocks.2.resnets.0.norm1.bias', None),
        ('decoder.up_blocks.2.resnets.0.conv1.weight', (2, 3, 1, 0)),
        ('decoder.up_blocks.2.resnets.0.conv1.bias', None),
        ('decoder.up_blocks.2.resnets.0.norm2.weight', None),
        ('decoder.up_blocks.2.resnets.0.norm2.bias', None),
        ('decoder.up_blocks.2.resnets.0.conv2.weight', (2, 3, 1, 0)),
        ('decoder.up_blocks.2.resnets.0.conv2.bias', None),
        ('decoder.up_blocks.2.resnets.0.conv_shortcut.weight', (2, 3, 1, 0)),
        ('decoder.up_blocks.2.resnets.0.conv_shortcut.bias', None),
        ('decoder.up_blocks.2.resnets.1.norm1.weight', None),
        ('decoder.up_blocks.2.resnets.1.norm1.bias', None),
        ('decoder.up_blocks.2.resnets.1.conv1.weight', (2, 3, 1, 0)),
        ('decoder.up_blocks.2.resnets.1.conv1.bias', None),
        ('decoder.up_blocks.2.resnets.1.norm2.weight', None),
        ('decoder.up_blocks.2.resnets.1.norm2.bias', None),
        ('decoder.up_blocks.2.resnets.1.conv2.weight', (2, 3, 1, 0)),
        ('decoder.up_blocks.2.resnets.1.conv2.bias', None),
        ('decoder.up_blocks.2.resnets.2.norm1.weight', None),
        ('decoder.up_blocks.2.resnets.2.norm1.bias', None),
        ('decoder.up_blocks.2.resnets.2.conv1.weight', (2, 3, 1, 0)),
        ('decoder.up_blocks.2.resnets.2.conv1.bias', None),
        ('decoder.up_blocks.2.resnets.2.norm2.weight', None),
        ('decoder.up_blocks.2.resnets.2.norm2.bias', None),
        ('decoder.up_blocks.2.resnets.2.conv2.weight', (2, 3, 1, 0)),
        ('decoder.up_blocks.2.resnets.2.conv2.bias', None),
        ('decoder.up_blocks.2.upsamplers.0.conv.weight', (2, 3, 1, 0)),
        ('decoder.up_blocks.2.upsamplers.0.conv.bias', None),
        ('decoder.up_blocks.3.resnets.0.norm1.weight', None),
        ('decoder.up_blocks.3.resnets.0.norm1.bias', None),
        ('decoder.up_blocks.3.resnets.0.conv1.weight', (2, 3, 1, 0)),
        ('decoder.up_blocks.3.resnets.0.conv1.bias', None),
        ('decoder.up_blocks.3.resnets.0.norm2.weight', None),
        ('decoder.up_blocks.3.resnets.0.norm2.bias', None),
        ('decoder.up_blocks.3.resnets.0.conv2.weight', (2, 3, 1, 0)),
        ('decoder.up_blocks.3.resnets.0.conv2.bias', None),
        ('decoder.up_blocks.3.resnets.0.conv_shortcut.weight', (2, 3, 1, 0)),
        ('decoder.up_blocks.3.resnets.0.conv_shortcut.bias', None),
        ('decoder.up_blocks.3.resnets.1.norm1.weight', None),
        ('decoder.up_blocks.3.resnets.1.norm1.bias', None),
        ('decoder.up_blocks.3.resnets.1.conv1.weight', (2, 3, 1, 0)),
        ('decoder.up_blocks.3.resnets.1.conv1.bias', None),
        ('decoder.up_blocks.3.resnets.1.norm2.weight', None),
        ('decoder.up_blocks.3.resnets.1.norm2.bias', None),
        ('decoder.up_blocks.3.resnets.1.conv2.weight', (2, 3, 1, 0)),
        ('decoder.up_blocks.3.resnets.1.conv2.bias', None),
        ('decoder.up_blocks.3.resnets.2.norm1.weight', None),
        ('decoder.up_blocks.3.resnets.2.norm1.bias', None),
        ('decoder.up_blocks.3.resnets.2.conv1.weight', (2, 3, 1, 0)),
        ('decoder.up_blocks.3.resnets.2.conv1.bias', None),
        ('decoder.up_blocks.3.resnets.2.norm2.weight', None),
        ('decoder.up_blocks.3.resnets.2.norm2.bias', None),
        ('decoder.up_blocks.3.resnets.2.conv2.weight', (2, 3, 1, 0)),
        ('decoder.up_blocks.3.resnets.2.conv2.bias', None),
        ('decoder.conv_norm_out.weight', None),
        ('decoder.conv_norm_out.bias', None),
        ('decoder.conv_out.weight', (2, 3, 1, 0)),
        ('decoder.conv_out.bias', None),
    ],
    'controlnet': [('control_model.time_embed.0.weight', (1, 0)),
                   ('control_model.time_embed.0.bias', None),
                   ('control_model.input_blocks.0.0.weight', (2, 3, 1, 0)),
                   ('control_model.input_blocks.0.0.bias', None),
                   ('control_model.time_embed.2.weight', (1, 0)),
                   ('control_model.time_embed.2.bias', None),
                   ('control_model.input_blocks.1.0.in_layers.0.weight', None),
                   ('control_model.input_blocks.1.0.in_layers.0.bias', None),
                   ('control_model.input_blocks.1.0.in_layers.2.weight', (2, 3, 1, 0)),
                   ('control_model.input_blocks.1.0.in_layers.2.bias', None),
                   ('control_model.input_blocks.1.0.emb_layers.1.weight', (1, 0)),
                   ('control_model.input_blocks.1.0.emb_layers.1.bias', None),
                   ('control_model.input_blocks.1.0.out_layers.0.weight', None),
                   ('control_model.input_blocks.1.0.out_layers.0.bias', None),
                   ('control_model.input_blocks.1.0.out_layers.3.weight', (2, 3, 1, 0)),
                   ('control_model.input_blocks.1.0.out_layers.3.bias', None),
                   ('control_model.input_blocks.1.1.norm.weight', None),
                   ('control_model.input_blocks.1.1.norm.bias', None),
                   ('control_model.input_blocks.1.1.proj_in.weight', (2, 3, 1, 0)),
                   ('control_model.input_blocks.1.1.proj_in.bias', None),
                   ('control_model.input_blocks.1.1.transformer_blocks.0.norm1.weight',
                    None),
                   ('control_model.input_blocks.1.1.transformer_blocks.0.norm1.bias',
                    None),
                   ('control_model.input_blocks.1.1.transformer_blocks.0.attn1.to_q.weight',
                    (1, 0)),
                   ('control_model.input_blocks.1.1.transformer_blocks.0.attn1.to_k.weight',
                    (1, 0)),
                   ('control_model.input_blocks.1.1.transformer_blocks.0.attn1.to_v.weight',
                    (1, 0)),
                   ('control_model.input_blocks.1.1.transformer_blocks.0.attn1.to_out.0.weight',
                    (1, 0)),
                   ('control_model.input_blocks.1.1.transformer_blocks.0.attn1.to_out.0.bias',
                    None),
                   ('control_model.input_blocks.1.1.transformer_blocks.0.norm2.weight',
                    None),
                   ('control_model.input_blocks.1.1.transformer_blocks.0.norm2.bias',
                    None),
                   ('control_model.input_blocks.1.1.transformer_blocks.0.attn2.to_q.weight',
                    (1, 0)),
                   ('control_model.input_blocks.1.1.transformer_blocks.0.attn2.to_k.weight',
                    (1, 0)),
                   ('control_model.input_blocks.1.1.transformer_blocks.0.attn2.to_v.weight',
                    (1, 0)),
                   ('control_model.input_blocks.1.1.transformer_blocks.0.attn2.to_out.0.weight',
                    (1, 0)),
                   ('control_model.input_blocks.1.1.transformer_blocks.0.attn2.to_out.0.bias',
                    None),
                   ('control_model.input_blocks.1.1.transformer_blocks.0.norm3.weight',
                    None),
                   ('control_model.input_blocks.1.1.transformer_blocks.0.norm3.bias',
                    None),
                   ('control_model.input_blocks.1.1.transformer_blocks.0.ff.net.0.proj.weight',
                    (1, 0)),
                   ('control_model.input_blocks.1.1.transformer_blocks.0.ff.net.0.proj.bias',
                    None),
                   ('control_model.input_blocks.1.1.transformer_blocks.0.ff.net.2.weight',
                    (1, 0)),
                   ('control_model.input_blocks.1.1.transformer_blocks.0.ff.net.2.bias',
                    None),
                   ('control_model.input_blocks.1.1.proj_out.weight', (2, 3, 1, 0)),
                   ('control_model.input_blocks.1.1.proj_out.bias', None),
                   ('control_model.input_blocks.2.0.in_layers.0.weight', None),
                   ('control_model.input_blocks.2.0.in_layers.0.bias', None),
                   ('control_model.input_blocks.2.0.in_layers.2.weight', (2, 3, 1, 0)),
                   ('control_model.input_blocks.2.0.in_layers.2.bias', None),
                   ('control_model.input_blocks.2.0.emb_layers.1.weight', (1, 0)),
                   ('control_model.input_blocks.2.0.emb_layers.1.bias', None),
                   ('control_model.input_blocks.2.0.out_layers.0.weight', None),
                   ('control_model.input_blocks.2.0.out_layers.0.bias', None),
                   ('control_model.input_blocks.2.0.out_layers.3.weight', (2, 3, 1, 0)),
                   ('control_model.input_blocks.2.0.out_layers.3.bias', None),
                   ('control_model.input_blocks.2.1.norm.weight', None),
                   ('control_model.input_blocks.2.1.norm.bias', None),
                   ('control_model.input_blocks.2.1.proj_in.weight', (2, 3, 1, 0)),
                   ('control_model.input_blocks.2.1.proj_in.bias', None),
                   ('control_model.input_blocks.2.1.transformer_blocks.0.norm1.weight',
                    None),
                   ('control_model.input_blocks.2.1.transformer_blocks.0.norm1.bias',
                    None),
                   ('control_model.input_blocks.2.1.transformer_blocks.0.attn1.to_q.weight',
                    (1, 0)),
                   ('control_model.input_blocks.2.1.transformer_blocks.0.attn1.to_k.weight',
                    (1, 0)),
                   ('control_model.input_blocks.2.1.transformer_blocks.0.attn1.to_v.weight',
                    (1, 0)),
                   ('control_model.input_blocks.2.1.transformer_blocks.0.attn1.to_out.0.weight',
                    (1, 0)),
                   ('control_model.input_blocks.2.1.transformer_blocks.0.attn1.to_out.0.bias',
                    None),
                   ('control_model.input_blocks.2.1.transformer_blocks.0.norm2.weight',
                    None),
                   ('control_model.input_blocks.2.1.transformer_blocks.0.norm2.bias',
                    None),
                   ('control_model.input_blocks.2.1.transformer_blocks.0.attn2.to_q.weight',
                    (1, 0)),
                   ('control_model.input_blocks.2.1.transformer_blocks.0.attn2.to_k.weight',
                    (1, 0)),
                   ('control_model.input_blocks.2.1.transformer_blocks.0.attn2.to_v.weight',
                    (1, 0)),
                   ('control_model.input_blocks.2.1.transformer_blocks.0.attn2.to_out.0.weight',
                    (1, 0)),
                   ('control_model.input_blocks.2.1.transformer_blocks.0.attn2.to_out.0.bias',
                    None),
                   ('control_model.input_blocks.2.1.transformer_blocks.0.norm3.weight',
                    None),
                   ('control_model.input_blocks.2.1.transformer_blocks.0.norm3.bias',
                    None),
                   ('control_model.input_blocks.2.1.transformer_blocks.0.ff.net.0.proj.weight',
                    (1, 0)),
                   ('control_model.input_blocks.2.1.transformer_blocks.0.ff.net.0.proj.bias',
                    None),
                   ('control_model.input_blocks.2.1.transformer_blocks.0.ff.net.2.weight',
                    (1, 0)),
                   ('control_model.input_blocks.2.1.transformer_blocks.0.ff.net.2.bias',
                    None),
                   ('control_model.input_blocks.2.1.proj_out.weight', (2, 3, 1, 0)),
                   ('control_model.input_blocks.2.1.proj_out.bias', None),
                   ('control_model.input_blocks.3.0.op.weight', (2, 3, 1, 0)),
                   ('control_model.input_blocks.3.0.op.bias', None),
                   ('control_model.input_blocks.4.0.in_layers.0.weight', None),
                   ('control_model.input_blocks.4.0.in_layers.0.bias', None),
                   ('control_model.input_blocks.4.0.in_layers.2.weight', (2, 3, 1, 0)),
                   ('control_model.input_blocks.4.0.in_layers.2.bias', None),
                   ('control_model.input_blocks.4.0.emb_layers.1.weight', (1, 0)),
                   ('control_model.input_blocks.4.0.emb_layers.1.bias', None),
                   ('control_model.input_blocks.4.0.out_layers.0.weight', None),
                   ('control_model.input_blocks.4.0.out_layers.0.bias', None),
                   ('control_model.input_blocks.4.0.out_layers.3.weight', (2, 3, 1, 0)),
                   ('control_model.input_blocks.4.0.out_layers.3.bias', None),
                   ('control_model.input_blocks.4.0.skip_connection.weight',
                    (2, 3, 1, 0)),
                   ('control_model.input_blocks.4.0.skip_connection.bias', None),
                   ('control_model.input_blocks.4.1.norm.weight', None),
                   ('control_model.input_blocks.4.1.norm.bias', None),
                   ('control_model.input_blocks.4.1.proj_in.weight', (2, 3, 1, 0)),
                   ('control_model.input_blocks.4.1.proj_in.bias', None),
                   ('control_model.input_blocks.4.1.transformer_blocks.0.norm1.weight',
                    None),
                   ('control_model.input_blocks.4.1.transformer_blocks.0.norm1.bias',
                    None),
                   ('control_model.input_blocks.4.1.transformer_blocks.0.attn1.to_q.weight',
                    (1, 0)),
                   ('control_model.input_blocks.4.1.transformer_blocks.0.attn1.to_k.weight',
                    (1, 0)),
                   ('control_model.input_blocks.4.1.transformer_blocks.0.attn1.to_v.weight',
                    (1, 0)),
                   ('control_model.input_blocks.4.1.transformer_blocks.0.attn1.to_out.0.weight',
                    (1, 0)),
                   ('control_model.input_blocks.4.1.transformer_blocks.0.attn1.to_out.0.bias',
                    None),
                   ('control_model.input_blocks.4.1.transformer_blocks.0.norm2.weight',
                    None),
                   ('control_model.input_blocks.4.1.transformer_blocks.0.norm2.bias',
                    None),
                   ('control_model.input_blocks.4.1.transformer_blocks.0.attn2.to_q.weight',
                    (1, 0)),
                   ('control_model.input_blocks.4.1.transformer_blocks.0.attn2.to_k.weight',
                    (1, 0)),
                   ('control_model.input_blocks.4.1.transformer_blocks.0.attn2.to_v.weight',
                    (1, 0)),
                   ('control_model.input_blocks.4.1.transformer_blocks.0.attn2.to_out.0.weight',
                    (1, 0)),
                   ('control_model.input_blocks.4.1.transformer_blocks.0.attn2.to_out.0.bias',
                    None),
                   ('control_model.input_blocks.4.1.transformer_blocks.0.norm3.weight',
                    None),
                   ('control_model.input_blocks.4.1.transformer_blocks.0.norm3.bias',
                    None),
                   ('control_model.input_blocks.4.1.transformer_blocks.0.ff.net.0.proj.weight',
                    (1, 0)),
                   ('control_model.input_blocks.4.1.transformer_blocks.0.ff.net.0.proj.bias',
                    None),
                   ('control_model.input_blocks.4.1.transformer_blocks.0.ff.net.2.weight',
                    (1, 0)),
                   ('control_model.input_blocks.4.1.transformer_blocks.0.ff.net.2.bias',
                    None),
                   ('control_model.input_blocks.4.1.proj_out.weight', (2, 3, 1, 0)),
                   ('control_model.input_blocks.4.1.proj_out.bias', None),
                   ('control_model.input_blocks.5.0.in_layers.0.weight', None),
                   ('control_model.input_blocks.5.0.in_layers.0.bias', None),
                   ('control_model.input_blocks.5.0.in_layers.2.weight', (2, 3, 1, 0)),
                   ('control_model.input_blocks.5.0.in_layers.2.bias', None),
                   ('control_model.input_blocks.5.0.emb_layers.1.weight', (1, 0)),
                   ('control_model.input_blocks.5.0.emb_layers.1.bias', None),
                   ('control_model.input_blocks.5.0.out_layers.0.weight', None),
                   ('control_model.input_blocks.5.0.out_layers.0.bias', None),
                   ('control_model.input_blocks.5.0.out_layers.3.weight', (2, 3, 1, 0)),
                   ('control_model.input_blocks.5.0.out_layers.3.bias', None),
                   ('control_model.input_blocks.5.1.norm.weight', None),
                   ('control_model.input_blocks.5.1.norm.bias', None),
                   ('control_model.input_blocks.5.1.proj_in.weight', (2, 3, 1, 0)),
                   ('control_model.input_blocks.5.1.proj_in.bias', None),
                   ('control_model.input_blocks.5.1.transformer_blocks.0.norm1.weight',
                    None),
                   ('control_model.input_blocks.5.1.transformer_blocks.0.norm1.bias',
                    None),
                   ('control_model.input_blocks.5.1.transformer_blocks.0.attn1.to_q.weight',
                    (1, 0)),
                   ('control_model.input_blocks.5.1.transformer_blocks.0.attn1.to_k.weight',
                    (1, 0)),
                   ('control_model.input_blocks.5.1.transformer_blocks.0.attn1.to_v.weight',
                    (1, 0)),
                   ('control_model.input_blocks.5.1.transformer_blocks.0.attn1.to_out.0.weight',
                    (1, 0)),
                   ('control_model.input_blocks.5.1.transformer_blocks.0.attn1.to_out.0.bias',
                    None),
                   ('control_model.input_blocks.5.1.transformer_blocks.0.norm2.weight',
                    None),
                   ('control_model.input_blocks.5.1.transformer_blocks.0.norm2.bias',
                    None),
                   ('control_model.input_blocks.5.1.transformer_blocks.0.attn2.to_q.weight',
                    (1, 0)),
                   ('control_model.input_blocks.5.1.transformer_blocks.0.attn2.to_k.weight',
                    (1, 0)),
                   ('control_model.input_blocks.5.1.transformer_blocks.0.attn2.to_v.weight',
                    (1, 0)),
                   ('control_model.input_blocks.5.1.transformer_blocks.0.attn2.to_out.0.weight',
                    (1, 0)),
                   ('control_model.input_blocks.5.1.transformer_blocks.0.attn2.to_out.0.bias',
                    None),
                   ('control_model.input_blocks.5.1.transformer_blocks.0.norm3.weight',
                    None),
                   ('control_model.input_blocks.5.1.transformer_blocks.0.norm3.bias',
                    None),
                   ('control_model.input_blocks.5.1.transformer_blocks.0.ff.net.0.proj.weight',
                    (1, 0)),
                   ('control_model.input_blocks.5.1.transformer_blocks.0.ff.net.0.proj.bias',
                    None),
                   ('control_model.input_blocks.5.1.transformer_blocks.0.ff.net.2.weight',
                    (1, 0)),
                   ('control_model.input_blocks.5.1.transformer_blocks.0.ff.net.2.bias',
                    None),
                   ('control_model.input_blocks.5.1.proj_out.weight', (2, 3, 1, 0)),
                   ('control_model.input_blocks.5.1.proj_out.bias', None),
                   ('control_model.input_blocks.6.0.op.weight', (2, 3, 1, 0)),
                   ('control_model.input_blocks.6.0.op.bias', None),
                   ('control_model.input_blocks.7.0.in_layers.0.weight', None),
                   ('control_model.input_blocks.7.0.in_layers.0.bias', None),
                   ('control_model.input_blocks.7.0.in_layers.2.weight', (2, 3, 1, 0)),
                   ('control_model.input_blocks.7.0.in_layers.2.bias', None),
                   ('control_model.input_blocks.7.0.emb_layers.1.weight', (1, 0)),
                   ('control_model.input_blocks.7.0.emb_layers.1.bias', None),
                   ('control_model.input_blocks.7.0.out_layers.0.weight', None),
                   ('control_model.input_blocks.7.0.out_layers.0.bias', None),
                   ('control_model.input_blocks.7.0.out_layers.3.weight', (2, 3, 1, 0)),
                   ('control_model.input_blocks.7.0.out_layers.3.bias', None),
                   ('control_model.input_blocks.7.0.skip_connection.weight',
                    (2, 3, 1, 0)),
                   ('control_model.input_blocks.7.0.skip_connection.bias', None),
                   ('control_model.input_blocks.7.1.norm.weight', None),
                   ('control_model.input_blocks.7.1.norm.bias', None),
                   ('control_model.input_blocks.7.1.proj_in.weight', (2, 3, 1, 0)),
                   ('control_model.input_blocks.7.1.proj_in.bias', None),
                   ('control_model.input_blocks.7.1.transformer_blocks.0.norm1.weight',
                    None),
                   ('control_model.input_blocks.7.1.transformer_blocks.0.norm1.bias',
                    None),
                   ('control_model.input_blocks.7.1.transformer_blocks.0.attn1.to_q.weight',
                    (1, 0)),
                   ('control_model.input_blocks.7.1.transformer_blocks.0.attn1.to_k.weight',
                    (1, 0)),
                   ('control_model.input_blocks.7.1.transformer_blocks.0.attn1.to_v.weight',
                    (1, 0)),
                   ('control_model.input_blocks.7.1.transformer_blocks.0.attn1.to_out.0.weight',
                    (1, 0)),
                   ('control_model.input_blocks.7.1.transformer_blocks.0.attn1.to_out.0.bias',
                    None),
                   ('control_model.input_blocks.7.1.transformer_blocks.0.norm2.weight',
                    None),
                   ('control_model.input_blocks.7.1.transformer_blocks.0.norm2.bias',
                    None),
                   ('control_model.input_blocks.7.1.transformer_blocks.0.attn2.to_q.weight',
                    (1, 0)),
                   ('control_model.input_blocks.7.1.transformer_blocks.0.attn2.to_k.weight',
                    (1, 0)),
                   ('control_model.input_blocks.7.1.transformer_blocks.0.attn2.to_v.weight',
                    (1, 0)),
                   ('control_model.input_blocks.7.1.transformer_blocks.0.attn2.to_out.0.weight',
                    (1, 0)),
                   ('control_model.input_blocks.7.1.transformer_blocks.0.attn2.to_out.0.bias',
                    None),
                   ('control_model.input_blocks.7.1.transformer_blocks.0.norm3.weight',
                    None),
                   ('control_model.input_blocks.7.1.transformer_blocks.0.norm3.bias',
                    None),
                   ('control_model.input_blocks.7.1.transformer_blocks.0.ff.net.0.proj.weight',
                    (1, 0)),
                   ('control_model.input_blocks.7.1.transformer_blocks.0.ff.net.0.proj.bias',
                    None),
                   ('control_model.input_blocks.7.1.transformer_blocks.0.ff.net.2.weight',
                    (1, 0)),
                   ('control_model.input_blocks.7.1.transformer_blocks.0.ff.net.2.bias',
                    None),
                   ('control_model.input_blocks.7.1.proj_out.weight', (2, 3, 1, 0)),
                   ('control_model.input_blocks.7.1.proj_out.bias', None),
                   ('control_model.input_blocks.8.0.in_layers.0.weight', None),
                   ('control_model.input_blocks.8.0.in_layers.0.bias', None),
                   ('control_model.input_blocks.8.0.in_layers.2.weight', (2, 3, 1, 0)),
                   ('control_model.input_blocks.8.0.in_layers.2.bias', None),
                   ('control_model.input_blocks.8.0.emb_layers.1.weight', (1, 0)),
                   ('control_model.input_blocks.8.0.emb_layers.1.bias', None),
                   ('control_model.input_blocks.8.0.out_layers.0.weight', None),
                   ('control_model.input_blocks.8.0.out_layers.0.bias', None),
                   ('control_model.input_blocks.8.0.out_layers.3.weight', (2, 3, 1, 0)),
                   ('control_model.input_blocks.8.0.out_layers.3.bias', None),
                   ('control_model.input_blocks.8.1.norm.weight', None),
                   ('control_model.input_blocks.8.1.norm.bias', None),
                   ('control_model.input_blocks.8.1.proj_in.weight', (2, 3, 1, 0)),
                   ('control_model.input_blocks.8.1.proj_in.bias', None),
                   ('control_model.input_blocks.8.1.transformer_blocks.0.norm1.weight',
                    None),
                   ('control_model.input_blocks.8.1.transformer_blocks.0.norm1.bias',
                    None),
                   ('control_model.input_blocks.8.1.transformer_blocks.0.attn1.to_q.weight',
                    (1, 0)),
                   ('control_model.input_blocks.8.1.transformer_blocks.0.attn1.to_k.weight',
                    (1, 0)),
                   ('control_model.input_blocks.8.1.transformer_blocks.0.attn1.to_v.weight',
                    (1, 0)),
                   ('control_model.input_blocks.8.1.transformer_blocks.0.attn1.to_out.0.weight',
                    (1, 0)),
                   ('control_model.input_blocks.8.1.transformer_blocks.0.attn1.to_out.0.bias',
                    None),
                   ('control_model.input_blocks.8.1.transformer_blocks.0.norm2.weight',
                    None),
                   ('control_model.input_blocks.8.1.transformer_blocks.0.norm2.bias',
                    None),
                   ('control_model.input_blocks.8.1.transformer_blocks.0.attn2.to_q.weight',
                    (1, 0)),
                   ('control_model.input_blocks.8.1.transformer_blocks.0.attn2.to_k.weight',
                    (1, 0)),
                   ('control_model.input_blocks.8.1.transformer_blocks.0.attn2.to_v.weight',
                    (1, 0)),
                   ('control_model.input_blocks.8.1.transformer_blocks.0.attn2.to_out.0.weight',
                    (1, 0)),
                   ('control_model.input_blocks.8.1.transformer_blocks.0.attn2.to_out.0.bias',
                    None),
                   ('control_model.input_blocks.8.1.transformer_blocks.0.norm3.weight',
                    None),
                   ('control_model.input_blocks.8.1.transformer_blocks.0.norm3.bias',
                    None),
                   ('control_model.input_blocks.8.1.transformer_blocks.0.ff.net.0.proj.weight',
                    (1, 0)),
                   ('control_model.input_blocks.8.1.transformer_blocks.0.ff.net.0.proj.bias',
                    None),
                   ('control_model.input_blocks.8.1.transformer_blocks.0.ff.net.2.weight',
                    (1, 0)),
                   ('control_model.input_blocks.8.1.transformer_blocks.0.ff.net.2.bias',
                    None),
                   ('control_model.input_blocks.8.1.proj_out.weight', (2, 3, 1, 0)),
                   ('control_model.input_blocks.8.1.proj_out.bias', None),
                   ('control_model.input_blocks.9.0.op.weight', (2, 3, 1, 0)),
                   ('control_model.input_blocks.9.0.op.bias', None),
                   ('control_model.input_blocks.10.0.in_layers.0.weight', None),
                   ('control_model.input_blocks.10.0.in_layers.0.bias', None),
                   ('control_model.input_blocks.10.0.in_layers.2.weight', (2, 3, 1, 0)),
                   ('control_model.input_blocks.10.0.in_layers.2.bias', None),
                   ('control_model.input_blocks.10.0.emb_layers.1.weight', (1, 0)),
                   ('control_model.input_blocks.10.0.emb_layers.1.bias', None),
                   ('control_model.input_blocks.10.0.out_layers.0.weight', None),
                   ('control_model.input_blocks.10.0.out_layers.0.bias', None),
                   ('control_model.input_blocks.10.0.out_layers.3.weight',
                    (2, 3, 1, 0)),
                   ('control_model.input_blocks.10.0.out_layers.3.bias', None),
                   ('control_model.input_blocks.11.0.in_layers.0.weight', None),
                   ('control_model.input_blocks.11.0.in_layers.0.bias', None),
                   ('control_model.input_blocks.11.0.in_layers.2.weight', (2, 3, 1, 0)),
                   ('control_model.input_blocks.11.0.in_layers.2.bias', None),
                   ('control_model.input_blocks.11.0.emb_layers.1.weight', (1, 0)),
                   ('control_model.input_blocks.11.0.emb_layers.1.bias', None),
                   ('control_model.input_blocks.11.0.out_layers.0.weight', None),
                   ('control_model.input_blocks.11.0.out_layers.0.bias', None),
                   ('control_model.input_blocks.11.0.out_layers.3.weight',
                    (2, 3, 1, 0)),
                   ('control_model.input_blocks.11.0.out_layers.3.bias', None),
                   ('control_model.middle_block.0.in_layers.0.weight', None),
                   ('control_model.middle_block.0.in_layers.0.bias', None),
                   ('control_model.middle_block.0.in_layers.2.weight', (2, 3, 1, 0)),
                   ('control_model.middle_block.0.in_layers.2.bias', None),
                   ('control_model.middle_block.0.emb_layers.1.weight', (1, 0)),
                   ('control_model.middle_block.0.emb_layers.1.bias', None),
                   ('control_model.middle_block.0.out_layers.0.weight', None),
                   ('control_model.middle_block.0.out_layers.0.bias', None),
                   ('control_model.middle_block.0.out_layers.3.weight', (2, 3, 1, 0)),
                   ('control_model.middle_block.0.out_layers.3.bias', None),
                   ('control_model.middle_block.1.norm.weight', None),
                   ('control_model.middle_block.1.norm.bias', None),
                   ('control_model.middle_block.1.proj_in.weight', (2, 3, 1, 0)),
                   ('control_model.middle_block.1.proj_in.bias', None),
                   ('control_model.middle_block.1.transformer_blocks.0.norm1.weight',
                    None),
                   ('control_model.middle_block.1.transformer_blocks.0.norm1.bias',
                    None),
                   ('control_model.middle_block.1.transformer_blocks.0.attn1.to_q.weight',
                    (1, 0)),
                   ('control_model.middle_block.1.transformer_blocks.0.attn1.to_k.weight',
                    (1, 0)),
                   ('control_model.middle_block.1.transformer_blocks.0.attn1.to_v.weight',
                    (1, 0)),
                   ('control_model.middle_block.1.transformer_blocks.0.attn1.to_out.0.weight',
                    (1, 0)),
                   ('control_model.middle_block.1.transformer_blocks.0.attn1.to_out.0.bias',
                    None),
                   ('control_model.middle_block.1.transformer_blocks.0.norm2.weight',
                    None),
                   ('control_model.middle_block.1.transformer_blocks.0.norm2.bias',
                    None),
                   ('control_model.middle_block.1.transformer_blocks.0.attn2.to_q.weight',
                    (1, 0)),
                   ('control_model.middle_block.1.transformer_blocks.0.attn2.to_k.weight',
                    (1, 0)),
                   ('control_model.middle_block.1.transformer_blocks.0.attn2.to_v.weight',
                    (1, 0)),
                   ('control_model.middle_block.1.transformer_blocks.0.attn2.to_out.0.weight',
                    (1, 0)),
                   ('control_model.middle_block.1.transformer_blocks.0.attn2.to_out.0.bias',
                    None),
                   ('control_model.middle_block.1.transformer_blocks.0.norm3.weight',
                    None),
                   ('control_model.middle_block.1.transformer_blocks.0.norm3.bias',
                    None),
                   ('control_model.middle_block.1.transformer_blocks.0.ff.net.0.proj.weight',
                    (1, 0)),
                   ('control_model.middle_block.1.transformer_blocks.0.ff.net.0.proj.bias',
                    None),
                   ('control_model.middle_block.1.transformer_blocks.0.ff.net.2.weight',
                    (1, 0)),
                   ('control_model.middle_block.1.transformer_blocks.0.ff.net.2.bias',
                    None),
                   ('control_model.middle_block.1.proj_out.weight', (2, 3, 1, 0)),
                   ('control_model.middle_block.1.proj_out.bias', None),
                   ('control_model.middle_block.2.in_layers.0.weight', None),
                   ('control_model.middle_block.2.in_layers.0.bias', None),
                   ('control_model.middle_block.2.in_layers.2.weight', (2, 3, 1, 0)),
                   ('control_model.middle_block.2.in_layers.2.bias', None),
                   ('control_model.middle_block.2.emb_layers.1.weight', (1, 0)),
                   ('control_model.middle_block.2.emb_layers.1.bias', None),
                   ('control_model.middle_block.2.out_layers.0.weight', None),
                   ('control_model.middle_block.2.out_layers.0.bias', None),
                   ('control_model.middle_block.2.out_layers.3.weight', (2, 3, 1, 0)),
                   ('control_model.middle_block.2.out_layers.3.bias', None),

                   ('control_model.zero_convs.0.0.weight', (2, 3, 1, 0)),
                   ('control_model.zero_convs.0.0.bias', None),
                   ('control_model.zero_convs.1.0.weight', (2, 3, 1, 0)),
                   ('control_model.zero_convs.1.0.bias', None),
                   ('control_model.zero_convs.2.0.weight', (2, 3, 1, 0)),
                   ('control_model.zero_convs.2.0.bias', None),
                   ('control_model.zero_convs.3.0.weight', (2, 3, 1, 0)),
                   ('control_model.zero_convs.3.0.bias', None),
                   ('control_model.zero_convs.4.0.weight', (2, 3, 1, 0)),
                   ('control_model.zero_convs.4.0.bias', None),
                   ('control_model.zero_convs.5.0.weight', (2, 3, 1, 0)),
                   ('control_model.zero_convs.5.0.bias', None),
                   ('control_model.zero_convs.6.0.weight', (2, 3, 1, 0)),
                   ('control_model.zero_convs.6.0.bias', None),
                   ('control_model.zero_convs.7.0.weight', (2, 3, 1, 0)),
                   ('control_model.zero_convs.7.0.bias', None),
                   ('control_model.zero_convs.8.0.weight', (2, 3, 1, 0)),
                   ('control_model.zero_convs.8.0.bias', None),
                   ('control_model.zero_convs.9.0.weight', (2, 3, 1, 0)),
                   ('control_model.zero_convs.9.0.bias', None),
                   ('control_model.zero_convs.10.0.weight', (2, 3, 1, 0)),
                   ('control_model.zero_convs.10.0.bias', None),
                   ('control_model.zero_convs.11.0.weight', (2, 3, 1, 0)),
                   ('control_model.zero_convs.11.0.bias', None),
                   ('control_model.middle_block_out.0.weight', (2, 3, 1, 0)),
                   ('control_model.middle_block_out.0.bias', None)],
    'hintnet': [
        ('control_model.input_hint_block.0.weight', (2, 3, 1, 0)),
        ('control_model.input_hint_block.0.bias', None),
        ('control_model.input_hint_block.2.weight', (2, 3, 1, 0)),
        ('control_model.input_hint_block.2.bias', None),
        ('control_model.input_hint_block.4.weight', (2, 3, 1, 0)),
        ('control_model.input_hint_block.4.bias', None),
        ('control_model.input_hint_block.6.weight', (2, 3, 1, 0)),
        ('control_model.input_hint_block.6.bias', None),
        ('control_model.input_hint_block.8.weight', (2, 3, 1, 0)),
        ('control_model.input_hint_block.8.bias', None),
        ('control_model.input_hint_block.10.weight', (2, 3, 1, 0)),
        ('control_model.input_hint_block.10.bias', None),
        ('control_model.input_hint_block.12.weight', (2, 3, 1, 0)),
        ('control_model.input_hint_block.12.bias', None),
        ('control_model.input_hint_block.14.weight', (2, 3, 1, 0)),
        ('control_model.input_hint_block.14.bias', None), ]
}


def load_weights_from_file(self, ckpt_path, ckpt_mapping, key_mapping=None, lora_dict=None):
    filename = os.path.basename(ckpt_path)
    print("{} loading:[{}]".format(self.name, filename))
    state_dict = {}
    if ckpt_path.endswith(".safetensors"):
        with safe_open(ckpt_path, framework="pt", device="cpu") as f:
            for key in f.keys():
                state_dict[key] = f.get_tensor(key)
    else:
        state_dict = torch.load(ckpt_path, map_location="cpu")
    print("loaded :[{}]".format(filename))
    weights = self.weights
    module_weights = []
    keys = state_dict.keys()
    if lora_dict is not None:
        lora_keys = lora_dict.keys()
    else:
        lora_keys = []
    lora_count = len(lora_keys)
    lora_idx = 0
    for i, (key, perm) in enumerate(ckpt_mapping):
        if isinstance(perm, str) and perm.endswith(".npy"):
            w = np.load(perm)
        else:
            if key_mapping is not None:
                if key in keys:
                    w = state_dict[key]
                else:
                    w = state_dict[key_mapping[key]]
            else:
                w = state_dict[key]
            if isinstance(w, torch.Tensor):
                w = w.detach().numpy()
            if lora_dict is not None:
                if key_mapping is not None:
                    lora_w = lora_dict.get(str(key_mapping.get(key, None)), None)
                else:
                    lora_w = lora_dict.get(key, None)
                if lora_w is not None:
                    w = w + lora_w
                    lora_idx += 1
            if perm is not None:
                w = np.transpose(w, perm)
        if weights[i].shape != w.shape:
            print("Wrong :[{},{}]".format(weights[i].name, key))
        module_weights.append(w)
    if lora_count > 0:
        if lora_idx == lora_count:
            print("Apply {} lora weights".format(lora_count))
        else:
            print("Apply {}/{} lora weights".format(lora_idx, lora_count))
    self.set_weights(module_weights)
    print("Loaded %d weights for %s" % (len(module_weights), self.name))


def load_weights_from_lora(ckpt_path):
    print("loading:[{}]".format(ckpt_path))
    state_dict = {}
    if ckpt_path.endswith(".safetensors"):
        with safe_open(ckpt_path, framework="pt", device="cpu") as f:
            for key in f.keys():
                state_dict[key] = f.get_tensor(key)
    else:
        state_dict = torch.load(ckpt_path, map_location="cpu")
    print("loaded :[{}]".format(ckpt_path))
    state_dict_keys = list(state_dict.keys())
    unet_state_dict = {}
    text_encoder_state_dict = {}
    print("try build lora dict:[{}]".format(ckpt_path))
    for idx, key in enumerate(state_dict_keys):
        if str(key).endswith(".alpha"):
            name = str(key).split(".alpha")[0]
            alpha = state_dict[name + ".alpha"]
            lora_down = state_dict[name + ".lora_down.weight"]
            lora_up = state_dict[name + ".lora_up.weight"]
            if isinstance(lora_down, torch.Tensor):
                lora_down = lora_down.detach().to(torch.float).cpu().numpy()
            if isinstance(lora_up, torch.Tensor):
                lora_up = lora_up.detach().to(torch.float).cpu().numpy()
            if isinstance(alpha, torch.Tensor):
                alpha = alpha.detach().to(torch.float).cpu().numpy()
            rank = float(lora_up.shape[1])
            w = np.einsum("ab...,bc...->ac...", lora_up, lora_down) * (alpha / rank)
            str_key = str(key)
            # todo : a more elegant implementation is needed
            if str_key.startswith("lora_te_text_model"):
                restore_name = name.replace("lora_te_text_model_encoder_layers_", '')
                restore_name = restore_name.replace("_mlp_fc1", '.mlp.fc1.weight')
                restore_name = restore_name.replace("_mlp_fc2", '.mlp.fc2.weight')
                restore_name = restore_name.replace("_self_attn_q_proj", '.self_attn.q_proj.weight')
                restore_name = restore_name.replace("_self_attn_v_proj", '.self_attn.v_proj.weight')
                restore_name = restore_name.replace("_self_attn_k_proj", '.self_attn.k_proj.weight')
                restore_name = restore_name.replace("_self_attn_out_proj", '.self_attn.out_proj.weight')
                restore_name = 'text_model.encoder.layers.' + restore_name
                text_encoder_state_dict[restore_name] = w
            elif str_key.startswith("lora_unet_"):
                restore_name = name.replace("lora_unet_", '')
                restore_name = restore_name.replace("down_blocks_", 'down_blocks.')
                restore_name = restore_name.replace("up_blocks_", 'up_blocks.')
                restore_name = restore_name.replace("_attentions_", '.attentions.')
                restore_name = restore_name.replace("_transformer_blocks_", '.transformer_blocks.')
                restore_name = restore_name.replace("_proj_in", '.proj_in.weight')
                restore_name = restore_name.replace("_proj_out", '.proj_out.weight')
                restore_name = restore_name.replace("_attn1_to_k", '.attn1.to_k.weight')
                restore_name = restore_name.replace("_attn2_to_k", '.attn2.to_k.weight')
                restore_name = restore_name.replace("_attn1_to_q", '.attn1.to_q.weight')
                restore_name = restore_name.replace("_attn2_to_q", '.attn2.to_q.weight')
                restore_name = restore_name.replace("_attn1_to_v", '.attn1.to_v.weight')
                restore_name = restore_name.replace("_attn2_to_v", '.attn2.to_v.weight')
                restore_name = restore_name.replace("_attn1_to_out_0", '.attn1.to_out.0.weight')
                restore_name = restore_name.replace("_attn2_to_out_0", '.attn2.to_out.0.weight')
                restore_name = restore_name.replace("_ff_net_0_proj", '.ff.net.0.proj.weight')
                restore_name = restore_name.replace("_ff_net_2", '.ff.net.2.weight')
                unet_state_dict[restore_name] = w
    print("lora dict:[{}] done.".format(ckpt_path))
    return text_encoder_state_dict, unet_state_dict

stable_diffusion/ckpt_loader.py
cpuimage/minSDTF
# Copyright 2022 The KerasCV Authors
#
# Licensed under the Apache License, Version 2.0 (the "License");
# you may not use this file except in compliance with the License.
# You may obtain a copy of the License at
#
#     https://www.apache.org/licenses/LICENSE-2.0
#
# Unless required by applicable law or agreed to in writing, software
# distributed under the License is distributed on an "AS IS" BASIS,
# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
# See the License for the specific language governing permissions and
# limitations under the License.
"""This code is taken nearly verbatim from
https://github.com/divamgupta/stable-diffusion-tensorflow."""

import gzip
import html
from functools import lru_cache

import regex as re
import tensorflow as tf


@lru_cache()
def bytes_to_unicode():
    """Return a list of utf-8 bytes and a corresponding list of unicode strings.

    The reversible bpe codes work on unicode strings.
    This means you need a large # of unicode characters in your vocab if you
    want to avoid UNKs. When you're at something like a 10B token dataset you
    end up needing around 5K for decent coverage. This is a significant
    percentage of your normal, say, 32K bpe vocab. To avoid that, we want
    lookup tables between utf-8 bytes and unicode strings.
    And avoids mapping to whitespace/control characters the bpe code barfs on.
    """
    bs = (
            list(range(ord("!"), ord("~") + 1))
            + list(range(ord("¡"), ord("¬") + 1))
            + list(range(ord("®"), ord("ÿ") + 1))
    )
    cs = bs[:]
    n = 0
    for b in range(2 ** 8):
        if b not in bs:
            bs.append(b)
            cs.append(2 ** 8 + n)
            n += 1
    cs = [chr(n) for n in cs]
    return dict(zip(bs, cs))


def get_pairs(word):
    """Return set of symbol pairs in a word.

    A word is represented as tuple of symbols(symbols being variable-length
    strings).
    """
    pairs = set()
    prev_char = word[0]
    for char in word[1:]:
        pairs.add((prev_char, char))
        prev_char = char
    return pairs


def basic_clean(text):
    text = html.unescape(html.unescape(text))
    return text.strip()


def whitespace_clean(text):
    text = re.sub(r"\s+", " ", text)
    text = text.strip()
    return text


class SimpleTokenizer:
    def __init__(self, bpe_path=None):
        bpe_path = bpe_path or tf.keras.utils.get_file(
            "bpe_simple_vocab_16e6.txt.gz",
            "https://github.com/openai/CLIP/blob/main/clip/bpe_simple_vocab_16e6.txt.gz?raw=true",  # noqa: E501
            file_hash="924691ac288e54409236115652ad4aa250f48203de50a9e4722a6ecd48d6804a",  # noqa: E501
        )
        self.byte_encoder = bytes_to_unicode()
        self.byte_decoder = {v: k for k, v in self.byte_encoder.items()}
        merges = gzip.open(bpe_path).read().decode("utf-8").split("\n")
        merges = merges[1: 49152 - 256 - 2 + 1]
        merges = [tuple(merge.split()) for merge in merges]
        vocab = list(bytes_to_unicode().values())
        vocab = vocab + [v + "</w>" for v in vocab]
        for merge in merges:
            vocab.append("".join(merge))
        vocab.extend(["<|startoftext|>", "<|endoftext|>"])
        self.vocab = vocab
        self.encoder = self._create_encoder(self.vocab)
        self.decoder = self._create_decoder(self.encoder)
        self.bpe_ranks = dict(zip(merges, range(len(merges))))

        self.special_tokens = {
            "<|startoftext|>": "<|startoftext|>",
            "<|endoftext|>": "<|endoftext|>",
        }
        self.cache = {
            "<|startoftext|>": "<|startoftext|>",
            "<|endoftext|>": "<|endoftext|>",
        }
        self.pat = self._create_pat()

    def _create_encoder(self, vocab):
        return dict(zip(vocab, range(len(vocab))))

    def _create_decoder(self, encoder):
        return {v: k for k, v in encoder.items()}

    def _create_pat(self):
        return re.compile(
            "|".join([re.escape(key) for key in self.special_tokens.keys()])
            + r"""|'s|'t|'re|'ve|'m|'ll|'d|[\p{L}]+|[\p{N}]|[^\s\p{L}\p{N}]+""",
            re.IGNORECASE,
        )

    @property
    def end_of_text(self):
        return self.encoder["<|endoftext|>"]

    @property
    def start_of_text(self):
        return self.encoder["<|startoftext|>"]

    def add_tokens(self, tokens):
        if isinstance(tokens, str):
            tokens = [tokens]
        tokens_added = 0
        for token in tokens:
            if token in self.vocab:
                continue
            tokens_added += 1
            self.vocab.append(token)
            self.special_tokens[token] = token
            self.cache[token] = token
        self.encoder = self._create_encoder(self.vocab)
        self.decoder = self._create_decoder(self.encoder)
        self.pat = self._create_pat()
        return tokens_added

    def bpe(self, token):
        if token in self.cache:
            return self.cache[token]
        word = tuple(token[:-1]) + (token[-1] + "</w>",)
        pairs = get_pairs(word)

        if not pairs:
            return token + "</w>"

        while True:
            bigram = min(
                pairs, key=lambda pair: self.bpe_ranks.get(pair, float("inf"))
            )
            if bigram not in self.bpe_ranks:
                break
            first, second = bigram
            new_word = []
            i = 0
            while i < len(word):
                try:
                    j = word.index(first, i)
                    new_word.extend(word[i:j])
                    i = j
                except:
                    new_word.extend(word[i:])
                    break

                if (word[i] == first
                        and i < len(word) - 1
                        and word[i + 1] == second):
                    new_word.append(first + second)
                    i += 2
                else:
                    new_word.append(word[i])
                    i += 1
            new_word = tuple(new_word)
            word = new_word
            if len(word) == 1:
                break
            else:
                pairs = get_pairs(word)
        word = " ".join(word)
        self.cache[token] = word
        return word

    def encode(self, text):
        bpe_tokens = []
        text = whitespace_clean(basic_clean(text)).lower()
        for token in re.findall(self.pat, text):
            token = "".join(self.byte_encoder[b] for b in token.encode("utf-8"))
            bpe_tokens.extend(
                self.encoder[bpe_token]
                for bpe_token in self.bpe(token).split(" ")
            )
        return [self.start_of_text] + bpe_tokens + [self.end_of_text]

    def decode(self, tokens):
        text = "".join([self.decoder[token] for token in tokens])
        text = (
            bytearray([self.byte_decoder[c] for c in text])
            .decode("utf-8", errors="replace")
            .replace("</w>", " ")
        )
        return text

stable_diffusion/clip_tokenizer.py
cpuimage/minSDTF
import os

import tensorflow as tf

from .ckpt_loader import load_weights_from_file, CKPT_MAPPING
from .diffusion_model import Attentions, ResBlock
from .layers import PaddedConv2D


class HintNet(tf.keras.Sequential):
    def __init__(self, img_height=512, img_width=512, name=None, controlnet_path=None):
        super().__init__(
            [
                tf.keras.layers.Input((img_height, img_width, 3)),
                PaddedConv2D(16, kernel_size=3, padding=1),
                tf.keras.layers.Activation("swish"),
                PaddedConv2D(16, kernel_size=3, padding=1),
                tf.keras.layers.Activation("swish"),
                PaddedConv2D(32, kernel_size=3, padding=1, strides=2),
                tf.keras.layers.Activation("swish"),
                PaddedConv2D(32, kernel_size=3, padding=1),
                tf.keras.layers.Activation("swish"),
                PaddedConv2D(96, kernel_size=3, padding=1, strides=2),
                tf.keras.layers.Activation("swish"),
                PaddedConv2D(96, kernel_size=3, padding=1),
                tf.keras.layers.Activation("swish"),
                PaddedConv2D(256, kernel_size=3, padding=1, strides=2),
                tf.keras.layers.Activation("swish"),
                PaddedConv2D(320, kernel_size=3, padding=1),
            ],
            name=name)
        origin = "https://huggingface.co/lllyasviel/ControlNet/resolve/main/models/control_sd15_canny.pth"
        ckpt_mapping = CKPT_MAPPING["hintnet"]
        if controlnet_path is not None:
            if os.path.exists(controlnet_path):
                load_weights_from_file(self, controlnet_path, ckpt_mapping=ckpt_mapping)
                return
            else:
                origin = controlnet_path
        model_weights_fpath = tf.keras.utils.get_file(origin=origin)
        if os.path.exists(model_weights_fpath):
            load_weights_from_file(self, model_weights_fpath, ckpt_mapping=ckpt_mapping)


class ControlNet(tf.keras.Model):
    def __init__(self, img_height=512, img_width=512, name=None, controlnet_path=None):
        context = tf.keras.layers.Input((None, 768))
        t_embed_input = tf.keras.layers.Input((320,))
        latent = tf.keras.layers.Input((img_height // 8, img_width // 8, 4))
        hint_out = tf.keras.layers.Input((img_height // 8, img_width // 8, 320))
        t_emb = tf.keras.layers.Dense(1280, name="time_embedding.linear_1")(t_embed_input)
        t_emb = tf.keras.layers.Activation("swish")(t_emb)
        t_emb = tf.keras.layers.Dense(1280, name="time_embedding.linear_2",
                                      activation=tf.keras.layers.Activation("swish"))(t_emb)
        outputs = []
        x = PaddedConv2D(320, kernel_size=3, padding=1)(latent) + hint_out
        outputs.append(x)
        x = ResBlock(320)([x, t_emb])
        x = Attentions(8, 40, fully_connected=False)([x, context])
        outputs.append(x)
        x = ResBlock(320)([x, t_emb])
        x = Attentions(8, 40, fully_connected=False)([x, context])
        outputs.append(x)
        x = PaddedConv2D(320, 3, strides=2, padding=1)(x)
        outputs.append(x)
        x = ResBlock(640)([x, t_emb])
        x = Attentions(8, 80, fully_connected=False)([x, context])
        outputs.append(x)
        x = ResBlock(640)([x, t_emb])
        x = Attentions(8, 80, fully_connected=False)([x, context])
        outputs.append(x)
        x = PaddedConv2D(640, 3, strides=2, padding=1)(x)
        outputs.append(x)
        x = ResBlock(1280)([x, t_emb])
        x = Attentions(8, 160, fully_connected=False)([x, context])
        outputs.append(x)
        x = ResBlock(1280)([x, t_emb])
        x = Attentions(8, 160, fully_connected=False)([x, context])
        outputs.append(x)
        x = PaddedConv2D(1280, 3, strides=2, padding=1)(x)
        outputs.append(x)
        x = ResBlock(1280)([x, t_emb])
        outputs.append(x)
        x = ResBlock(1280)([x, t_emb])
        outputs.append(x)
        x = ResBlock(1280)([x, t_emb])
        x = Attentions(8, 160, fully_connected=False)([x, context])
        x = ResBlock(1280)([x, t_emb])
        outputs.append(x)
        assert len(outputs) == 13
        outs = []
        zero_convs = [PaddedConv2D(320, kernel_size=1, padding=0),
                      PaddedConv2D(320, kernel_size=1, padding=0),
                      PaddedConv2D(320, kernel_size=1, padding=0),
                      PaddedConv2D(320, kernel_size=1, padding=0),
                      PaddedConv2D(640, kernel_size=1, padding=0),
                      PaddedConv2D(640, kernel_size=1, padding=0),
                      PaddedConv2D(640, kernel_size=1, padding=0),
                      PaddedConv2D(1280, kernel_size=1, padding=0),
                      PaddedConv2D(1280, kernel_size=1, padding=0),
                      PaddedConv2D(1280, kernel_size=1, padding=0),
                      PaddedConv2D(1280, kernel_size=1, padding=0),
                      PaddedConv2D(1280, kernel_size=1, padding=0),
                      PaddedConv2D(1280, kernel_size=1, padding=0)]
        for i, x in enumerate(outputs):
            outs.append(zero_convs[i](x))
        super().__init__([latent, t_embed_input, context, hint_out], outs, name=name)
        origin = "https://huggingface.co/lllyasviel/ControlNet/resolve/main/models/control_sd15_canny.pth"
        ckpt_mapping = CKPT_MAPPING["controlnet"]
        if controlnet_path is not None:
            if os.path.exists(controlnet_path):
                load_weights_from_file(self, controlnet_path, ckpt_mapping=ckpt_mapping)
                return
            else:
                origin = controlnet_path
        model_weights_fpath = tf.keras.utils.get_file(origin=origin)
        if os.path.exists(model_weights_fpath):
            load_weights_from_file(self, model_weights_fpath, ckpt_mapping=ckpt_mapping)

stable_diffusion/control_net.py
cpuimage/minSDTF
# Copyright 2022 The KerasCV Authors
#
# Licensed under the Apache License, Version 2.0 (the "License");
# you may not use this file except in compliance with the License.
# You may obtain a copy of the License at
#
#     https://www.apache.org/licenses/LICENSE-2.0
#
# Unless required by applicable law or agreed to in writing, software
# distributed under the License is distributed on an "AS IS" BASIS,
# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
# See the License for the specific language governing permissions and
# limitations under the License.
import os

import tensorflow as tf

from .ckpt_loader import load_weights_from_file, CKPT_MAPPING, UNET_KEY_MAPPING
from .layers import GroupNormalization, PaddedConv2D


class ResBlock(tf.keras.layers.Layer):
    def __init__(self, output_dim, **kwargs):
        super().__init__(**kwargs)
        self.output_dim = output_dim
        self.entry_flow = [
            GroupNormalization(epsilon=1e-5, name="norm1"),
            tf.keras.layers.Activation("swish"),
            PaddedConv2D(output_dim, 3, padding=1, name="conv1")]
        self.embedding_flow = tf.keras.layers.Dense(output_dim, name="time_emb_proj")
        self.exit_flow = [
            GroupNormalization(epsilon=1e-5, name="norm2"),
            tf.keras.layers.Activation("swish"),
            PaddedConv2D(output_dim, 3, padding=1, name="conv2")]

    def build(self, input_shape):
        if input_shape[0][-1] != self.output_dim:
            self.residual_projection = PaddedConv2D(self.output_dim, 1, name="conv_shortcut")
        else:
            self.residual_projection = lambda x: x

    def call(self, inputs):
        inputs, embeddings = inputs
        x = inputs
        for layer in self.entry_flow:
            x = layer(x)
        embeddings = self.embedding_flow(embeddings)
        x = x + embeddings[:, None, None]
        for layer in self.exit_flow:
            x = layer(x)
        return x + self.residual_projection(inputs)


class Attentions(tf.keras.layers.Layer):
    def __init__(self, num_heads, head_size, fully_connected=False, **kwargs):
        super().__init__(**kwargs)
        self.norm = GroupNormalization(epsilon=1e-5, name="norm")
        channels = num_heads * head_size
        if fully_connected:
            self.proj_in = tf.keras.layers.Dense(num_heads * head_size, name="proj_in")
        else:
            self.proj_in = PaddedConv2D(num_heads * head_size, 1, name="proj_in")
        self.transformer_block = TransformerBlock(channels, num_heads, head_size, name="transformer_blocks.0")
        if fully_connected:
            self.proj_out = tf.keras.layers.Dense(channels, name="proj_out")
        else:
            self.proj_out = PaddedConv2D(channels, 1, name="proj_out")

    def call(self, inputs):
        inputs, context = inputs
        _, h, w, c = inputs.shape
        x = self.norm(inputs)
        x = self.proj_in(x)
        x = tf.reshape(x, (-1, h * w, c))
        x = self.transformer_block([x, context])
        x = tf.reshape(x, (-1, h, w, c))
        return self.proj_out(x) + inputs


class TransformerBlock(tf.keras.layers.Layer):
    def __init__(self, dim, num_heads, head_size, **kwargs):
        super().__init__(**kwargs)
        self.norm1 = tf.keras.layers.LayerNormalization(epsilon=1e-5, name="norm1")
        self.attn1 = CrossAttention(num_heads, head_size, name="attn1")
        self.norm2 = tf.keras.layers.LayerNormalization(epsilon=1e-5, name="norm2")
        self.attn2 = CrossAttention(num_heads, head_size, name="attn2")
        self.norm3 = tf.keras.layers.LayerNormalization(epsilon=1e-5, name="norm3")
        self.geglu = GEGLU(dim * 4, name="ff.net.0")
        self.dense = tf.keras.layers.Dense(dim, name="ff.net.2")

    def call(self, inputs):
        inputs, context = inputs
        x = self.attn1([self.norm1(inputs), None]) + inputs
        x = self.attn2([self.norm2(x), context]) + x
        return self.dense(self.geglu(self.norm3(x))) + x


class CrossAttention(tf.keras.layers.Layer):
    def __init__(self, num_heads, head_size, **kwargs):
        super().__init__(**kwargs)
        self.to_q = tf.keras.layers.Dense(num_heads * head_size, use_bias=False, name="to_q")
        self.to_k = tf.keras.layers.Dense(num_heads * head_size, use_bias=False, name="to_k")
        self.to_v = tf.keras.layers.Dense(num_heads * head_size, use_bias=False, name="to_v")
        self.scale = head_size ** -0.5
        self.num_heads = num_heads
        self.head_size = head_size
        self.out_proj = tf.keras.layers.Dense(num_heads * head_size, name="to_out")

    def call(self, inputs):
        inputs, context = inputs
        context = inputs if context is None else context
        q, k, v = self.to_q(inputs), self.to_k(context), self.to_v(context)
        batch_size = tf.shape(inputs)[0]
        q = tf.reshape(q, (batch_size, inputs.shape[1], self.num_heads, self.head_size))
        k = tf.reshape(k, (batch_size, -1, self.num_heads, self.head_size))
        v = tf.reshape(v, (batch_size, -1, self.num_heads, self.head_size))

        q = tf.transpose(q, (0, 2, 1, 3))  # (bs, num_heads, time, head_size)
        k = tf.transpose(k, (0, 2, 3, 1))  # (bs, num_heads, head_size, time)
        v = tf.transpose(v, (0, 2, 1, 3))  # (bs, num_heads, time, head_size)

        # score = td_dot(q, k) * self.scale
        score = tf.einsum('bnqh,bnhk->bnqk', q, k) * self.scale
        weights = tf.keras.activations.softmax(score)  # (bs, num_heads, time, time)
        # attn = td_dot(weights, v)
        attn = tf.einsum('bnqk,bnkh->bnqh', weights, v)
        attn = tf.transpose(attn, (0, 2, 1, 3))  # (bs, time, num_heads, head_size)
        out = tf.reshape(attn, (-1, inputs.shape[1], self.num_heads * self.head_size))
        return self.out_proj(out)


class Upsamplers(tf.keras.layers.Layer):
    def __init__(self, channels, **kwargs):
        super().__init__(**kwargs)
        self.ups = tf.keras.layers.UpSampling2D(2)
        self.conv = PaddedConv2D(channels, 3, padding=1, name="conv")

    def call(self, inputs):
        return self.conv(self.ups(inputs))


class GEGLU(tf.keras.layers.Layer):
    def __init__(self, output_dim, **kwargs):
        super().__init__(**kwargs)
        self.output_dim = output_dim
        self.proj = tf.keras.layers.Dense(output_dim * 2, name="proj")

    def call(self, inputs):
        x = self.proj(inputs)
        x, gate = x[..., : self.output_dim], x[..., self.output_dim:]
        tanh_res = tf.keras.activations.tanh(
            gate * 0.7978845608 * (1 + 0.044715 * (gate ** 2)))
        return x * 0.5 * gate * (1 + tanh_res)


def td_dot(a, b):
    aa = tf.reshape(a, (-1, a.shape[2], a.shape[3]))
    bb = tf.reshape(b, (-1, b.shape[2], b.shape[3]))
    cc = tf.keras.backend.batch_dot(aa, bb)
    return tf.reshape(cc, (-1, a.shape[1], cc.shape[1], cc.shape[2]))


class DiffusionModel(tf.keras.Model):
    def __init__(self, img_height=512, img_width=512, apply_control_net=False, name=None,
                 ckpt_path=None, lora_dict=None):
        context = tf.keras.layers.Input((None, 768))
        t_embed_input = tf.keras.layers.Input((320,))
        latent = tf.keras.layers.Input((img_height // 8, img_width // 8, 4))
        controls = None
        if apply_control_net:
            controls = [tf.keras.layers.Input((img_height // 8, img_width // 8, 320)),
                        tf.keras.layers.Input((img_height // 8, img_width // 8, 320)),
                        tf.keras.layers.Input((img_height // 8, img_width // 8, 320)),
                        tf.keras.layers.Input((img_height // 16, img_width // 16, 320)),
                        tf.keras.layers.Input((img_height // 16, img_width // 16, 640)),
                        tf.keras.layers.Input((img_height // 16, img_width // 16, 640)),
                        tf.keras.layers.Input((img_height // 32, img_width // 32, 640)),
                        tf.keras.layers.Input((img_height // 32, img_width // 32, 1280)),
                        tf.keras.layers.Input((img_height // 32, img_width // 32, 1280)),
                        tf.keras.layers.Input((img_height // 64, img_width // 64, 1280)),
                        tf.keras.layers.Input((img_height // 64, img_width // 64, 1280)),
                        tf.keras.layers.Input((img_height // 64, img_width // 64, 1280)),
                        tf.keras.layers.Input((img_height // 64, img_width // 64, 1280))]
        t_emb = tf.keras.layers.Dense(1280, name="time_embedding.linear_1")(t_embed_input)
        t_emb = tf.keras.layers.Activation("swish")(t_emb)
        t_emb = tf.keras.layers.Dense(1280, name="time_embedding.linear_2",
                                      activation=tf.keras.layers.Activation("swish"))(
            t_emb)
        # Downsampling flow
        outputs = []
        x = PaddedConv2D(320, kernel_size=3, padding=1, name="conv_in")(latent)
        outputs.append(x)
        # down_blocks.0
        x = ResBlock(320, name="down_blocks.0.resnets.0")([x, t_emb])
        x = Attentions(8, 40, fully_connected=False, name="down_blocks.0.attentions.0")([x, context])
        outputs.append(x)
        x = ResBlock(320, name="down_blocks.0.resnets.1")([x, t_emb])
        x = Attentions(8, 40, fully_connected=False, name="down_blocks.0.attentions.1")([x, context])
        outputs.append(x)
        x = PaddedConv2D(320, 3, strides=2, padding=1, name="down_blocks.0.downsamplers.0")(x)  # Downsample 2x
        outputs.append(x)
        # down_blocks.1
        x = ResBlock(640, name="down_blocks.1.resnets.0")([x, t_emb])
        x = Attentions(8, 80, fully_connected=False, name="down_blocks.1.attentions.0")([x, context])
        outputs.append(x)
        x = ResBlock(640, name="down_blocks.1.resnets.1")([x, t_emb])
        x = Attentions(8, 80, fully_connected=False, name="down_blocks.1.attentions.1")([x, context])
        outputs.append(x)
        x = PaddedConv2D(640, 3, strides=2, padding=1, name="down_blocks.1.downsamplers.0")(x)  # Downsample 2x
        outputs.append(x)
        # down_blocks.2
        x = ResBlock(1280, name="down_blocks.2.resnets.0")([x, t_emb])
        x = Attentions(8, 160, fully_connected=False, name="down_blocks.2.attentions.0")([x, context])
        outputs.append(x)
        x = ResBlock(1280, name="down_blocks.2.resnets.1")([x, t_emb])
        x = Attentions(8, 160, fully_connected=False, name="down_blocks.2.attentions.1")([x, context])
        outputs.append(x)
        x = PaddedConv2D(1280, 3, strides=2, padding=1, name="down_blocks.2.downsamplers.0")(x)  # Downsample 2x
        outputs.append(x)
        # down_blocks.3
        x = ResBlock(1280, name="down_blocks.3.resnets.0")([x, t_emb])
        outputs.append(x)
        x = ResBlock(1280, name="down_blocks.3.resnets.1")([x, t_emb])
        outputs.append(x)
        # mid_block
        # Middle flow
        x = ResBlock(1280, name="mid_block.resnets.0")([x, t_emb])
        x = Attentions(8, 160, fully_connected=False, name="mid_block.attentions.0")([x, context])
        x = ResBlock(1280, name="mid_block.resnets.1")([x, t_emb])
        if controls is not None:
            x = x + controls[12]
            assert len(outputs) == 12
            for i in range(len(outputs)):
                outputs[i] = outputs[i] + controls[i]
        # Upsampling flow
        # up_blocks.0
        x = tf.keras.layers.Concatenate(axis=-1)([x, outputs.pop()])
        x = ResBlock(1280, name="up_blocks.0.resnets.0")([x, t_emb])
        x = tf.keras.layers.Concatenate(axis=-1)([x, outputs.pop()])
        x = ResBlock(1280, name="up_blocks.0.resnets.1")([x, t_emb])
        x = tf.keras.layers.Concatenate(axis=-1)([x, outputs.pop()])
        x = ResBlock(1280, name="up_blocks.0.resnets.2")([x, t_emb])
        x = Upsamplers(1280, name="up_blocks.0.upsamplers.0")(x)
        # up_blocks.1
        x = tf.keras.layers.Concatenate(axis=-1)([x, outputs.pop()])
        x = ResBlock(1280, name="up_blocks.1.resnets.0")([x, t_emb])
        x = Attentions(8, 160, fully_connected=False, name="up_blocks.1.attentions.0")([x, context])
        x = tf.keras.layers.Concatenate(axis=-1)([x, outputs.pop()])
        x = ResBlock(1280, name="up_blocks.1.resnets.1")([x, t_emb])
        x = Attentions(8, 160, fully_connected=False, name="up_blocks.1.attentions.1")([x, context])
        x = tf.keras.layers.Concatenate(axis=-1)([x, outputs.pop()])
        x = ResBlock(1280, name="up_blocks.1.resnets.2")([x, t_emb])
        x = Attentions(8, 160, fully_connected=False, name="up_blocks.1.attentions.2")([x, context])
        x = Upsamplers(1280, name="up_blocks.1.upsamplers.0")(x)
        # up_blocks.2
        x = tf.keras.layers.Concatenate(axis=-1)([x, outputs.pop()])
        x = ResBlock(640, name="up_blocks.2.resnets.0")([x, t_emb])
        x = Attentions(8, 80, fully_connected=False, name="up_blocks.2.attentions.0")([x, context])
        x = tf.keras.layers.Concatenate(axis=-1)([x, outputs.pop()])
        x = ResBlock(640, name="up_blocks.2.resnets.1")([x, t_emb])
        x = Attentions(8, 80, fully_connected=False, name="up_blocks.2.attentions.1")([x, context])
        x = tf.keras.layers.Concatenate(axis=-1)([x, outputs.pop()])
        x = ResBlock(640, name="up_blocks.2.resnets.2")([x, t_emb])
        x = Attentions(8, 80, fully_connected=False, name="up_blocks.2.attentions.2")([x, context])
        x = Upsamplers(640, name="up_blocks.2.upsamplers.0")(x)
        # up_blocks.3
        x = tf.keras.layers.Concatenate(axis=-1)([x, outputs.pop()])
        x = ResBlock(320, name="up_blocks.3.resnets.0")([x, t_emb])
        x = Attentions(8, 40, fully_connected=False, name="up_blocks.3.attentions.0")([x, context])
        x = tf.keras.layers.Concatenate(axis=-1)([x, outputs.pop()])
        x = ResBlock(320, name="up_blocks.3.resnets.1")([x, t_emb])
        x = Attentions(8, 40, fully_connected=False, name="up_blocks.3.attentions.1")([x, context])
        x = tf.keras.layers.Concatenate(axis=-1)([x, outputs.pop()])
        x = ResBlock(320, name="up_blocks.3.resnets.2")([x, t_emb])
        x = Attentions(8, 40, fully_connected=False, name="up_blocks.3.attentions.2")([x, context])
        # Exit flow
        x = GroupNormalization(epsilon=1e-5, name="conv_norm_out")(x)
        x = tf.keras.layers.Activation("swish")(x)
        output = PaddedConv2D(4, kernel_size=3, padding=1, name="conv_out")(x)
        if controls is not None:
            super().__init__([latent, t_embed_input, context] + list(controls), output, name=name)
        else:
            super().__init__([latent, t_embed_input, context], output, name=name)
        origin = "https://huggingface.co/dreamlike-art/dreamlike-photoreal-2.0/resolve/main/dreamlike-photoreal-2.0.safetensors"
        ckpt_mapping = CKPT_MAPPING["civitai_model"]
        if ckpt_path is not None:
            if os.path.exists(ckpt_path):
                load_weights_from_file(self, ckpt_path, ckpt_mapping=ckpt_mapping, key_mapping=UNET_KEY_MAPPING,
                                       lora_dict=lora_dict)
                return
            else:
                origin = ckpt_path
        model_weights_fpath = tf.keras.utils.get_file(origin=origin)
        if os.path.exists(model_weights_fpath):
            load_weights_from_file(self, model_weights_fpath, ckpt_mapping=ckpt_mapping, key_mapping=UNET_KEY_MAPPING,
                                   lora_dict=lora_dict)

stable_diffusion/diffusion_model.py
cpuimage/minSDTF
# Copyright 2022 The KerasCV Authors
#
# Licensed under the Apache License, Version 2.0 (the "License");
# you may not use this file except in compliance with the License.
# You may obtain a copy of the License at
#
#     https://www.apache.org/licenses/LICENSE-2.0
#
# Unless required by applicable law or agreed to in writing, software
# distributed under the License is distributed on an "AS IS" BASIS,
# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
# See the License for the specific language governing permissions and
# limitations under the License.
import os

import tensorflow as tf

from .layers import GroupNormalization, AttentionBlock, PaddedConv2D, ResnetBlock
from .ckpt_loader import load_weights_from_file, CKPT_MAPPING


class ImageDecoder(tf.keras.Sequential):
    def __init__(self, img_height, img_width, name=None, ckpt_path=None):
        super().__init__(
            [
                tf.keras.layers.Input((img_height // 8, img_width // 8, 4)),
                tf.keras.layers.Rescaling(1.0 / 0.18215),
                PaddedConv2D(4, 1),
                PaddedConv2D(512, 3, padding=1),
                ResnetBlock(512),
                AttentionBlock(512),
                ResnetBlock(512),
                ResnetBlock(512),
                ResnetBlock(512),
                ResnetBlock(512),
                tf.keras.layers.UpSampling2D(2),
                PaddedConv2D(512, 3, padding=1),
                ResnetBlock(512),
                ResnetBlock(512),
                ResnetBlock(512),
                tf.keras.layers.UpSampling2D(2),
                PaddedConv2D(512, 3, padding=1),
                ResnetBlock(256),
                ResnetBlock(256),
                ResnetBlock(256),
                tf.keras.layers.UpSampling2D(2),
                PaddedConv2D(256, 3, padding=1),
                ResnetBlock(128),
                ResnetBlock(128),
                ResnetBlock(128),
                GroupNormalization(epsilon=1e-5),
                tf.keras.layers.Activation("swish"),
                PaddedConv2D(3, 3, padding=1),
            ],
            name=name)
        origin = "https://huggingface.co/stabilityai/sd-vae-ft-mse/resolve/main/diffusion_pytorch_model.safetensors"
        ckpt_mapping = CKPT_MAPPING["decoder"]
        if ckpt_path is not None:
            if os.path.exists(ckpt_path):
                load_weights_from_file(self, ckpt_path, ckpt_mapping=ckpt_mapping)
                return
            else:
                origin = ckpt_path
        model_weights_fpath = tf.keras.utils.get_file(origin=origin)
        if os.path.exists(model_weights_fpath):
            load_weights_from_file(self, model_weights_fpath, ckpt_mapping=ckpt_mapping)

stable_diffusion/image_decoder.py
cpuimage/minSDTF
# Copyright 2022 The KerasCV Authors
#
# Licensed under the Apache License, Version 2.0 (the "License");
# you may not use this file except in compliance with the License.
# You may obtain a copy of the License at
#
#     https://www.apache.org/licenses/LICENSE-2.0
#
# Unless required by applicable law or agreed to in writing, software
# distributed under the License is distributed on an "AS IS" BASIS,
# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
# See the License for the specific language governing permissions and
# limitations under the License.
import os

import tensorflow as tf

from .layers import GroupNormalization, AttentionBlock, PaddedConv2D, ResnetBlock
from .ckpt_loader import load_weights_from_file, CKPT_MAPPING


class ImageEncoder(tf.keras.Sequential):
    """ImageEncoder is the VAE Encoder for StableDiffusion."""

    def __init__(self, ckpt_path=None):
        super().__init__(
            [
                tf.keras.layers.Input((None, None, 3)),
                PaddedConv2D(128, 3, padding=1),
                ResnetBlock(128),
                ResnetBlock(128),
                PaddedConv2D(128, 3, padding=((0, 1), (0, 1)), strides=2),
                ResnetBlock(256),
                ResnetBlock(256),
                PaddedConv2D(256, 3, padding=((0, 1), (0, 1)), strides=2),
                ResnetBlock(512),
                ResnetBlock(512),
                PaddedConv2D(512, 3, padding=((0, 1), (0, 1)), strides=2),
                ResnetBlock(512),
                ResnetBlock(512),
                ResnetBlock(512),
                AttentionBlock(512),
                ResnetBlock(512),
                GroupNormalization(epsilon=1e-5),
                tf.keras.layers.Activation("swish"),
                PaddedConv2D(8, 3, padding=1),
                PaddedConv2D(8, 1),
                tf.keras.layers.Lambda(lambda x: tf.split(x, num_or_size_splits=2, axis=-1)[0] * 0.18215),
            ])
        origin = "https://huggingface.co/stabilityai/sd-vae-ft-mse/resolve/main/diffusion_pytorch_model.safetensors"
        ckpt_mapping = CKPT_MAPPING["encoder"]
        if ckpt_path is not None:
            if os.path.exists(ckpt_path):
                load_weights_from_file(self, ckpt_path, ckpt_mapping=ckpt_mapping)
                return
            else:
                origin = ckpt_path
        model_weights_fpath = tf.keras.utils.get_file(origin=origin)
        if os.path.exists(model_weights_fpath):
            load_weights_from_file(self, model_weights_fpath, ckpt_mapping=ckpt_mapping)

stable_diffusion/image_encoder.py
cpuimage/minSDTF
# Copyright 2022 The KerasCV Authors
#
# Licensed under the Apache License, Version 2.0 (the "License");
# you may not use this file except in compliance with the License.
# You may obtain a copy of the License at
#
#     https://www.apache.org/licenses/LICENSE-2.0
#
# Unless required by applicable law or agreed to in writing, software
# distributed under the License is distributed on an "AS IS" BASIS,
# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
# See the License for the specific language governing permissions and
# limitations under the License.
import tensorflow as tf


class GroupNormalization(tf.keras.layers.Layer):
    """Group normalization layer.

    Group Normalization divides the channels into groups and computes
    within each group the mean and variance for normalization.
    Empirically, its accuracy is more stable than batch norm in a wide
    range of small batch sizes, if learning rate is adjusted linearly
    with batch sizes.

    Relation to Layer Normalization:
    If the number of groups is set to 1, then this operation becomes nearly
    identical to Layer Normalization (see Layer Normalization docs for details).

    Relation to Instance Normalization:
    If the number of groups is set to the input dimension (number of groups is
    equal to number of channels), then this operation becomes identical to
    Instance Normalization.

    Args:
      groups: Integer, the number of groups for Group Normalization. Can be in
        the range [1, N] where N is the input dimension. The input dimension
        must be divisible by the number of groups. Defaults to 32.
      axis: Integer or List/Tuple. The axis or axes to normalize across.
        Typically this is the features axis/axes. The left-out axes are
        typically the batch axis/axes. This argument defaults to `-1`, the last
        dimension in the input.
      epsilon: Small float added to variance to avoid dividing by zero. Defaults
        to 1e-3
      center: If True, add offset of `beta` to normalized tensor. If False,
        `beta` is ignored. Defaults to True.
      scale: If True, multiply by `gamma`. If False, `gamma` is not used.
        Defaults to True. When the next layer is linear (also e.g. `nn.relu`),
        this can be disabled since the scaling will be done by the next layer.
      beta_initializer: Initializer for the beta weight. Defaults to zeros.
      gamma_initializer: Initializer for the gamma weight. Defaults to ones.
      beta_regularizer: Optional regularizer for the beta weight. None by
        default.
      gamma_regularizer: Optional regularizer for the gamma weight. None by
        default.
      beta_constraint: Optional constraint for the beta weight. None by default.
      gamma_constraint: Optional constraint for the gamma weight. None by
        default.  Input shape: Arbitrary. Use the keyword argument `input_shape`
        (tuple of integers, does not include the samples axis) when using this
        layer as the first layer in a model.  Output shape: Same shape as input.
    Reference: - [Yuxin Wu & Kaiming He, 2018](https://arxiv.org/abs/1803.08494)
    """

    def __init__(
            self,
            groups=32,
            axis=-1,
            epsilon=1e-3,
            center=True,
            scale=True,
            beta_initializer="zeros",
            gamma_initializer="ones",
            beta_regularizer=None,
            gamma_regularizer=None,
            beta_constraint=None,
            gamma_constraint=None,
            **kwargs,
    ):
        super().__init__(**kwargs)
        self.supports_masking = True
        self.groups = groups
        self.axis = axis
        self.epsilon = epsilon
        self.center = center
        self.scale = scale
        self.beta_initializer = tf.keras.initializers.get(beta_initializer)
        self.gamma_initializer = tf.keras.initializers.get(gamma_initializer)
        self.beta_regularizer = tf.keras.regularizers.get(beta_regularizer)
        self.gamma_regularizer = tf.keras.regularizers.get(gamma_regularizer)
        self.beta_constraint = tf.keras.constraints.get(beta_constraint)
        self.gamma_constraint = tf.keras.constraints.get(gamma_constraint)

    def validate_axis(self, axis, input_shape):
        """Validate an axis value and returns its standardized form.

        Args:
          axis: Value to validate. Can be an integer or a list/tuple of integers.
            Integers may be negative.
          input_shape: Reference input shape that the axis/axes refer to.

        Returns:
          Normalized form of `axis`, i.e. a list with all-positive values.
        """
        input_shape = tf.TensorShape(input_shape)
        rank = input_shape.rank
        if not rank:
            raise ValueError(
                f"Input has undefined rank. Received: input_shape={input_shape}")
        # Convert axis to list and resolve negatives
        if isinstance(axis, int):
            axis = [axis]
        else:
            axis = list(axis)
        for idx, x in enumerate(axis):
            if x < 0:
                axis[idx] = rank + x
        # Validate axes
        for x in axis:
            if x < 0 or x >= rank:
                raise ValueError(
                    "Invalid value for `axis` argument. "
                    "Expected 0 <= axis < inputs.rank (with "
                    f"inputs.rank={rank}). Received: axis={tuple(axis)}")
        if len(axis) != len(set(axis)):
            raise ValueError(f"Duplicate axis: {tuple(axis)}")
        return axis

    def build(self, input_shape):
        self.validate_axis(self.axis, input_shape)
        dim = input_shape[self.axis]
        if dim is None:
            raise ValueError(
                f"Axis {self.axis} of input tensor should have a defined "
                "dimension but the layer received an input with shape "
                f"{input_shape}.")

        if self.groups == -1:
            self.groups = dim

        if dim < self.groups:
            raise ValueError(
                f"Number of groups ({self.groups}) cannot be more than the "
                f"number of channels ({dim}).")

        if dim % self.groups != 0:
            raise ValueError(
                f"Number of groups ({self.groups}) must be a multiple "
                f"of the number of channels ({dim}).")

        self.input_spec = tf.keras.layers.InputSpec(
            ndim=len(input_shape), axes={self.axis: dim})

        if self.scale:
            self.gamma = self.add_weight(
                shape=(dim,),
                name="gamma",
                initializer=self.gamma_initializer,
                regularizer=self.gamma_regularizer,
                constraint=self.gamma_constraint, )
        else:
            self.gamma = None

        if self.center:
            self.beta = self.add_weight(
                shape=(dim,),
                name="beta",
                initializer=self.beta_initializer,
                regularizer=self.beta_regularizer,
                constraint=self.beta_constraint, )
        else:
            self.beta = None

        super().build(input_shape)

    def call(self, inputs):
        input_shape = tf.shape(inputs)
        reshaped_inputs = self._reshape_into_groups(inputs)
        normalized_inputs = self._apply_normalization(
            reshaped_inputs, input_shape)
        return tf.reshape(normalized_inputs, input_shape)

    def _reshape_into_groups(self, inputs):
        input_shape = tf.shape(inputs)
        group_shape = [input_shape[i] for i in range(inputs.shape.rank)]
        group_shape[self.axis] = input_shape[self.axis] // self.groups
        group_shape.insert(self.axis, self.groups)
        group_shape = tf.stack(group_shape)
        reshaped_inputs = tf.reshape(inputs, group_shape)
        return reshaped_inputs

    def _apply_normalization(self, reshaped_inputs, input_shape):
        group_reduction_axes = list(range(1, reshaped_inputs.shape.rank))
        axis = -2 if self.axis == -1 else self.axis - 1
        group_reduction_axes.pop(axis)
        mean, variance = tf.nn.moments(
            reshaped_inputs, group_reduction_axes, keepdims=True)
        gamma, beta = self._get_reshaped_weights(input_shape)
        normalized_inputs = tf.nn.batch_normalization(
            reshaped_inputs,
            mean=mean,
            variance=variance,
            scale=gamma,
            offset=beta,
            variance_epsilon=self.epsilon, )
        return normalized_inputs

    def _get_reshaped_weights(self, input_shape):
        broadcast_shape = self._create_broadcast_shape(input_shape)
        gamma = None
        beta = None
        if self.scale:
            gamma = tf.reshape(self.gamma, broadcast_shape)
        if self.center:
            beta = tf.reshape(self.beta, broadcast_shape)
        return gamma, beta

    def _create_broadcast_shape(self, input_shape):
        broadcast_shape = [1] * tf.keras.backend.int_shape(input_shape)[0]
        broadcast_shape[self.axis] = input_shape[self.axis] // self.groups
        broadcast_shape.insert(self.axis, self.groups)
        return broadcast_shape

    def compute_output_shape(self, input_shape):
        return input_shape

    def get_config(self):
        config = {
            "groups": self.groups,
            "axis": self.axis,
            "epsilon": self.epsilon,
            "center": self.center,
            "scale": self.scale,
            "beta_initializer": tf.keras.initializers.serialize(self.beta_initializer),
            "gamma_initializer": tf.keras.initializers.serialize(self.gamma_initializer),
            "beta_regularizer": tf.keras.regularizers.serialize(self.beta_regularizer),
            "gamma_regularizer": tf.keras.regularizers.serialize(self.gamma_regularizer),
            "beta_constraint": tf.keras.constraints.serialize(self.beta_constraint),
            "gamma_constraint": tf.keras.constraints.serialize(self.gamma_constraint),
        }
        base_config = super().get_config()
        return {**base_config, **config}


class PaddedConv2D(tf.keras.layers.Layer):
    def __init__(self, filters, kernel_size, padding=0, strides=1, **kwargs):
        super().__init__(**kwargs)
        self.padding2d = tf.keras.layers.ZeroPadding2D(padding)
        self.conv2d = tf.keras.layers.Conv2D(filters, kernel_size, strides=strides)

    def call(self, inputs):
        x = self.padding2d(inputs)
        return self.conv2d(x)


class AttentionBlock(tf.keras.layers.Layer):
    def __init__(self, output_dim, **kwargs):
        super().__init__(**kwargs)
        self.output_dim = output_dim
        self.norm = GroupNormalization(epsilon=1e-5)
        self.q = tf.keras.layers.Dense(output_dim, use_bias=True, )
        self.k = tf.keras.layers.Dense(output_dim, use_bias=True, )
        self.v = tf.keras.layers.Dense(output_dim, use_bias=True, )
        self.proj_out = tf.keras.layers.Dense(output_dim, use_bias=True, )

    def call(self, inputs):
        x = self.norm(inputs)
        q, k, v = self.q(x), self.k(x), self.v(x)

        # Compute attention
        shape = tf.shape(q)
        h, w, c = shape[1], shape[2], shape[3]
        q = tf.reshape(q, (-1, h * w, c))  # b, hw, c
        k = tf.transpose(k, (0, 3, 1, 2))
        k = tf.reshape(k, (-1, c, h * w))  # b, c, hw
        y = q @ k
        y = y * 1 / tf.sqrt(tf.cast(c, self.compute_dtype))
        y = tf.keras.activations.softmax(y)

        # Attend to values
        v = tf.transpose(v, (0, 3, 1, 2))
        v = tf.reshape(v, (-1, c, h * w))
        y = tf.transpose(y, (0, 2, 1))
        x = v @ y
        x = tf.transpose(x, (0, 2, 1))
        x = tf.reshape(x, (-1, h, w, c))
        return self.proj_out(x) + inputs


class ResnetBlock(tf.keras.layers.Layer):
    def __init__(self, output_dim, **kwargs):
        super().__init__(**kwargs)
        self.output_dim = output_dim
        self.norm1 = GroupNormalization(epsilon=1e-5)
        self.conv1 = PaddedConv2D(output_dim, 3, padding=1)
        self.norm2 = GroupNormalization(epsilon=1e-5)
        self.conv2 = PaddedConv2D(output_dim, 3, padding=1)

    def build(self, input_shape):
        if input_shape[-1] != self.output_dim:
            self.residual_projection = PaddedConv2D(self.output_dim, 1)
        else:
            self.residual_projection = lambda x: x

    def call(self, inputs):
        x = self.conv1(tf.keras.activations.swish(self.norm1(inputs)))
        x = self.conv2(tf.keras.activations.swish(self.norm2(x)))
        return x + self.residual_projection(inputs)

stable_diffusion/layers.py
cpuimage/minSDTF
import re
from typing import List, Optional, Union

import numpy as np

re_attention = re.compile(
    r"""
\\\(|
\\\)|
\\\[|
\\]|
\\\\|
\\|
\(|
\[|
:([+-]?[.\d]+)\)|
\)|
]|
[^\\()\[\]:]+|
:
""",
    re.X,
)


def parse_prompt_attention(text):
    """
    Parses a string with attention tokens and returns a list of pairs: text and its associated weight.
    Accepted tokens are:
      (abc) - increases attention to abc by a multiplier of 1.1
      (abc:3.12) - increases attention to abc by a multiplier of 3.12
      [abc] - decreases attention to abc by a multiplier of 1.1
      \( - literal character '('
      \[ - literal character '['
      \) - literal character ')'
      \] - literal character ']'
      \\ - literal character '\'
      anything else - just text
    >>> parse_prompt_attention('normal text')
    [['normal text', 1.0]]
    >>> parse_prompt_attention('an (important) word')
    [['an ', 1.0], ['important', 1.1], [' word', 1.0]]
    >>> parse_prompt_attention('(unbalanced')
    [['unbalanced', 1.1]]
    >>> parse_prompt_attention('\(literal\]')
    [['(literal]', 1.0]]
    >>> parse_prompt_attention('(unnecessary)(parens)')
    [['unnecessaryparens', 1.1]]
    >>> parse_prompt_attention('a (((house:1.3)) [on] a (hill:0.5), sun, (((sky))).')
    [['a ', 1.0],
     ['house', 1.5730000000000004],
     [' ', 1.1],
     ['on', 1.0],
     [' a ', 1.1],
     ['hill', 0.55],
     [', sun, ', 1.1],
     ['sky', 1.4641000000000006],
     ['.', 1.1]]
    """

    res = []
    round_brackets = []
    square_brackets = []

    round_bracket_multiplier = 1.1
    square_bracket_multiplier = 1 / 1.1

    def multiply_range(start_position, multiplier):
        for p in range(start_position, len(res)):
            res[p][1] *= multiplier

    for m in re_attention.finditer(text):
        text = m.group(0)
        weight = m.group(1)

        if text.startswith("\\"):
            res.append([text[1:], 1.0])
        elif text == "(":
            round_brackets.append(len(res))
        elif text == "[":
            square_brackets.append(len(res))
        elif weight is not None and len(round_brackets) > 0:
            multiply_range(round_brackets.pop(), float(weight))
        elif text == ")" and len(round_brackets) > 0:
            multiply_range(round_brackets.pop(), round_bracket_multiplier)
        elif text == "]" and len(square_brackets) > 0:
            multiply_range(square_brackets.pop(), square_bracket_multiplier)
        else:
            res.append([text, 1.0])

    for pos in round_brackets:
        multiply_range(pos, round_bracket_multiplier)

    for pos in square_brackets:
        multiply_range(pos, square_bracket_multiplier)

    if len(res) == 0:
        res = [["", 1.0]]

    # merge runs of identical weights
    i = 0
    while i + 1 < len(res):
        if res[i][1] == res[i + 1][1]:
            res[i][0] += res[i + 1][0]
            res.pop(i + 1)
        else:
            i += 1

    return res


def get_prompts_with_weights(tokenizer, prompt: List[str], max_length: int, embedding_tokens_count: int = 0,
                             embedding_tokens_weight: float = 1.0):
    r"""
    Tokenize a list of prompts and return its tokens with weights of each token.

    No padding, starting or ending token is included.
    """
    tokens = []
    weights = []
    truncated = False
    for text in prompt:
        texts_and_weights = parse_prompt_attention(text)
        text_token = []
        text_weight = []
        if embedding_tokens_count > 0:
            # add embedding_tokens for init
            embedding_tokens = tokenizer.encode("*")[1:-1]
            text_token += embedding_tokens * embedding_tokens_count
            text_weight += [embedding_tokens_weight] * embedding_tokens_count
        for word, weight in texts_and_weights:
            # tokenize and discard the starting and the ending token
            token = tokenizer.encode(word.strip())[1:-1]
            text_token += list(token)
            # copy the weight by length of token
            text_weight += [weight] * len(token)
            # stop if the text is too long (longer than truncation limit)
            if len(text_token) > max_length:
                truncated = True
                break
        # truncate
        if len(text_token) > max_length:
            truncated = True
            text_token = text_token[:max_length]
            text_weight = text_weight[:max_length]
        tokens.append(text_token)
        weights.append(text_weight)
    if truncated:
        print("Prompt was truncated. Try to shorten the prompt or increase max_embeddings_multiples")
    return tokens, weights


def pad_tokens_and_weights(tokens, weights, max_length, bos, eos, pad, no_boseos_middle=True, chunk_length=77):
    r"""
    Pad the tokens (with starting and ending tokens) and weights (with 1.0) to max_length.
    """
    max_embeddings_multiples = (max_length - 2) // (chunk_length - 2)
    weights_length = max_length if no_boseos_middle else max_embeddings_multiples * chunk_length
    for i in range(len(tokens)):
        tokens[i] = [bos] + tokens[i] + [pad] * (max_length - 1 - len(tokens[i]) - 1) + [eos]
        if no_boseos_middle:
            weights[i] = [1.0] + weights[i] + [1.0] * (max_length - 1 - len(weights[i]))
        else:
            w = []
            if len(weights[i]) == 0:
                w = [1.0] * weights_length
            else:
                for j in range(max_embeddings_multiples):
                    w.append(1.0)  # weight for starting token in this chunk
                    w += weights[i][j * (chunk_length - 2): min(len(weights[i]), (j + 1) * (chunk_length - 2))]
                    w.append(1.0)  # weight for ending token in this chunk
                w += [1.0] * (weights_length - len(w))
            weights[i] = w[:]

    return tokens, weights


def get_unweighted_text_embeddings(
        text_clip_embedding,
        text_encoder,
        text_input: np.array,
        chunk_length: int,
        no_boseos_middle: Optional[bool] = True,
        embedding_tokens_count=0,
        embedding=None,
):
    """
    When the length of tokens is a multiple of the capacity of the text encoder,
    it should be split into chunks and sent to the text encoder individually.
    """
    use_embedding = embedding_tokens_count > 0 and embedding is not None
    max_embeddings_multiples = (text_input.shape[1] - 2) // (chunk_length - 2)
    if max_embeddings_multiples > 1:
        text_embeddings = []
        for i in range(max_embeddings_multiples):
            # extract the i-th chunk
            text_input_chunk = text_input[:, i * (chunk_length - 2): (i + 1) * (chunk_length - 2) + 2].copy()

            # cover the head and the tail by the starting and the ending tokens
            text_input_chunk[:, 0] = text_input[0, 0]
            text_input_chunk[:, -1] = text_input[0, -1]
            if use_embedding and i == 0:
                clip_embedding = text_clip_embedding.predict_on_batch(
                    [text_input_chunk, np.asarray([list(range(text_input_chunk.shape[1]))], dtype=np.int32)])
                clip_embedding = np.concatenate(
                    [clip_embedding[:, 0:1, :],
                     np.tile(embedding, (clip_embedding.shape[0], 1, 1)).astype(clip_embedding.dtype),
                     clip_embedding[:, embedding_tokens_count + 1:, :]], axis=1)
                text_embedding = text_encoder.predict_on_batch(clip_embedding)
            else:
                clip_embedding = text_clip_embedding.predict_on_batch(
                    [text_input_chunk, np.asarray([list(range(text_input_chunk.shape[1]))], dtype=np.int32)])
                text_embedding = text_encoder.predict_on_batch(clip_embedding)

            if no_boseos_middle:
                if i == 0:
                    # discard the ending token
                    text_embedding = text_embedding[:, :-1]
                elif i == max_embeddings_multiples - 1:
                    # discard the starting token
                    text_embedding = text_embedding[:, 1:]
                else:
                    # discard both starting and ending tokens
                    text_embedding = text_embedding[:, 1:-1]

            text_embeddings.append(text_embedding)
        text_embeddings = np.concatenate(text_embeddings, axis=1)
    else:
        clip_embedding = text_clip_embedding.predict_on_batch(
            [text_input, np.asarray([list(range(text_input.shape[1]))], dtype=np.int32)])
        if use_embedding:
            clip_embedding = np.concatenate(
                [clip_embedding[:, 0:1, :],
                 np.tile(embedding, (clip_embedding.shape[0], 1, 1)).astype(clip_embedding.dtype),
                 clip_embedding[:, embedding_tokens_count + 1:, :]], axis=1)
        text_embeddings = text_encoder.predict_on_batch(clip_embedding)
    return text_embeddings


def get_weighted_text_embeddings(
        tokenizer,
        text_clip_embedding,
        text_encoder,
        prompt: Union[str, List[str]],
        max_embeddings_multiples: Optional[int] = 4,
        no_boseos_middle: Optional[bool] = False,
        skip_parsing: Optional[bool] = False,
        skip_weighting: Optional[bool] = False,
        model_max_length=77,
        pad_token_id=49407,
        embedding_tokens_count=0,
        embedding_tokens_weight=1.0,
        embedding=None,
):
    r"""
    Prompts can be assigned with local weights using brackets. For example,
    prompt 'A (very beautiful) masterpiece' highlights the words 'very beautiful',
    and the embedding tokens corresponding to the words get multiplied by a constant, 1.1.

    Also, to regularize of the embedding, the weighted embedding would be scaled to preserve the original mean.

    Args:
        tokenizer  : provide access to the tokenizer
        text_encoder :  provide access to the text encoder.
        prompt (`str` or `List[str]`):
            The prompt or prompts to guide the image generation.
        max_embeddings_multiples (`int`, *optional*, defaults to `1`):
            The max multiple length of prompt embeddings compared to the max output length of text encoder.
        no_boseos_middle (`bool`, *optional*, defaults to `False`):
            If the length of text token is multiples of the capacity of text encoder, whether reserve the starting and
            ending token in each of the chunk in the middle.
        skip_parsing (`bool`, *optional*, defaults to `False`):
            Skip the parsing of brackets.
        skip_weighting (`bool`, *optional*, defaults to `False`):
            Skip the weighting. When the parsing is skipped, it is forced True.
    """
    if embedding_tokens_count > 0 and embedding is None:
        embedding_tokens_count = 0
    max_length = (model_max_length - 2) * max_embeddings_multiples + 2
    if isinstance(prompt, str):
        prompt = [prompt]

    if not skip_parsing:
        prompt_tokens, prompt_weights = get_prompts_with_weights(tokenizer, prompt, max_length - 2,
                                                                 embedding_tokens_count, embedding_tokens_weight)
    else:
        prompt_tokens = [
            token[1:-1]
            for token in tokenizer.encode(prompt)[:max_length]
        ]
        prompt_weights = [[1.0] * len(token) for token in prompt_tokens]

    # round up the longest length of tokens to a multiple of (model_max_length - 2)
    max_length = max([len(token) for token in prompt_tokens])

    max_embeddings_multiples = min(
        max_embeddings_multiples,
        (max_length - 1) // (model_max_length - 2) + 1,
    )
    max_embeddings_multiples = max(1, max_embeddings_multiples)
    max_length = (model_max_length - 2) * max_embeddings_multiples + 2

    # pad the length of tokens and weights
    bos = tokenizer.start_of_text
    eos = tokenizer.end_of_text
    pad = pad_token_id
    prompt_tokens, prompt_weights = pad_tokens_and_weights(
        prompt_tokens,
        prompt_weights,
        max_length,
        bos,
        eos,
        pad,
        no_boseos_middle=no_boseos_middle,
        chunk_length=model_max_length,
    )
    prompt_tokens = np.array(prompt_tokens, dtype=np.int32)
    # get the embeddings
    text_embeddings = get_unweighted_text_embeddings(
        text_clip_embedding,
        text_encoder,
        prompt_tokens,
        model_max_length,
        no_boseos_middle=no_boseos_middle,
        embedding_tokens_count=embedding_tokens_count,
        embedding=embedding,
    )
    prompt_weights = np.array(prompt_weights, dtype=text_embeddings.dtype)
    if (not skip_parsing) and (not skip_weighting):
        previous_mean = text_embeddings.mean(axis=(-2, -1))
        text_embeddings *= prompt_weights[:, :, None]
        text_embeddings *= (previous_mean / text_embeddings.mean(axis=(-2, -1)))[:, None, None]
    return text_embeddings

stable_diffusion/long_prompt_weighting.py
cpuimage/minSDTF
# Copyright 2022 The KerasCV Authors
#
# Licensed under the Apache License, Version 2.0 (the "License");
# you may not use this file except in compliance with the License.
# You may obtain a copy of the License at
#
#     https://www.apache.org/licenses/LICENSE-2.0
#
# Unless required by applicable law or agreed to in writing, software
# distributed under the License is distributed on an "AS IS" BASIS,
# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
# See the License for the specific language governing permissions and
# limitations under the License.
"""Keras implementation of StableDiffusion.

Credits:

- Original implementation:
  https://github.com/CompVis/stable-diffusion
- Initial TF/Keras port:
  https://github.com/divamgupta/stable-diffusion-tensorflow

The current implementation is a rewrite of the initial TF/Keras port by
Divam Gupta.
"""
import os

import numpy as np
import tensorflow as tf
import torch
from PIL import Image
from scipy.ndimage import correlate1d

from .ckpt_loader import load_weights_from_lora
from .clip_tokenizer import SimpleTokenizer
from .control_net import ControlNet, HintNet
from .diffusion_model import DiffusionModel
from .image_decoder import ImageDecoder
from .image_encoder import ImageEncoder
from .long_prompt_weighting import get_weighted_text_embeddings
from .text_encoder import TextEncoder, TextClipEmbedding

MAX_PROMPT_LENGTH = 77


class StableDiffusionBase:
    """Base class for stable diffusion and stable diffusion v2 model."""

    def __init__(
            self,
            img_height=512,
            img_width=512,
            jit_compile=False):
        # UNet requires multiples of 2**7 = 128
        img_height = round(img_height / 128) * 128
        img_width = round(img_width / 128) * 128
        self.img_height = img_height
        self.img_width = img_width

        # lazy initialize the component models and the tokenizer
        self._image_encoder = None
        self._text_encoder = None
        self._hint_net = None
        self._control_net = None
        self._text_clip_embedding = None
        self._diffusion_model = None
        self._image_decoder = None
        self._tokenizer = None
        self.jit_compile = jit_compile
        self.alphas_cumprod = np.cumprod(1. - np.square(np.linspace(np.sqrt(0.00085), np.sqrt(0.012), 1000)), axis=0)
        self.signal_rates = np.sqrt(self.alphas_cumprod)
        self.noise_rates = np.sqrt(1. - self.alphas_cumprod)

    def load_embedding(self, embedding_path):
        embedding = None
        if os.path.exists(str(embedding_path)):
            state_dict = torch.load(embedding_path, map_location="cpu")
            key_str = "string_to_param"
            if key_str in state_dict.keys():
                string_to_param = state_dict[key_str]
                for i, key in enumerate(list(string_to_param.keys())):
                    value = string_to_param[key]
                    if value.dtype == torch.float32 or value.dtype == torch.float16:
                        embedding = value.detach().numpy()
        return embedding

    def text_to_image(
            self,
            prompt,
            negative_prompt=None,
            batch_size=1,
            num_steps=50,
            unconditional_guidance_scale=7.5,
            embedding=None,
            negative_embedding=None,
            seed=None,
            control_net_image=None):
        encoded_text = self.encode_text(prompt, embedding)

        return self.generate_image(
            encoded_text,
            negative_prompt=negative_prompt,
            batch_size=batch_size,
            num_steps=num_steps,
            unconditional_guidance_scale=unconditional_guidance_scale,
            seed=seed,
            negative_embedding=negative_embedding,
            control_net_image=control_net_image)

    def image_to_image(
            self,
            prompt,
            negative_prompt=None,
            batch_size=1,
            num_steps=50,
            unconditional_guidance_scale=7.5,
            embedding=None,
            negative_embedding=None,
            seed=None,
            control_net_image=None,
            reference_image=None,
            reference_image_strength=0.8):
        encoded_text = self.encode_text(prompt, embedding)
        return self.generate_image(
            encoded_text,
            negative_prompt=negative_prompt,
            batch_size=batch_size,
            num_steps=num_steps,
            unconditional_guidance_scale=unconditional_guidance_scale,
            seed=seed,
            negative_embedding=negative_embedding,
            control_net_image=control_net_image,
            reference_image=reference_image,
            reference_image_strength=reference_image_strength)

    def inpaint(
            self,
            prompt,
            negative_prompt=None,
            batch_size=1,
            num_steps=50,
            unconditional_guidance_scale=7.5,
            embedding=None,
            negative_embedding=None,
            seed=None,
            control_net_image=None,
            reference_image=None,
            reference_image_strength=0.8,
            inpaint_mask=None,
            mask_blur_strength=5):
        encoded_text = self.encode_text(prompt, embedding)

        return self.generate_image(
            encoded_text,
            negative_prompt=negative_prompt,
            batch_size=batch_size,
            num_steps=num_steps,
            unconditional_guidance_scale=unconditional_guidance_scale,
            seed=seed,
            negative_embedding=negative_embedding,
            control_net_image=control_net_image,
            reference_image=reference_image,
            reference_image_strength=reference_image_strength,
            inpaint_mask=inpaint_mask,
            mask_blur_strength=mask_blur_strength)

    def encode_text(self, prompt, embedding_data=None):
        """Encodes a prompt into a latent text encoding.

        The encoding produced by this method should be used as the
        `encoded_text` parameter of `StableDiffusion.generate_image`. Encoding
        text separately from generating an image can be used to arbitrarily
        modify the text encoding prior to image generation, e.g. for walking
        between two prompts.

        Args:
            prompt: a string to encode, must be 77 tokens or shorter.
            embedding_data : for apply textual inversion
        Example:

        ```python
        from keras_cv.models import StableDiffusion

        model = StableDiffusion(img_height=512, img_width=512, jit_compile=True)
        encoded_text  = model.encode_text("Tacos at dawn")
        img = model.generate_image(encoded_text)
        ```
        """
        # Tokenize prompt (i.e. starting context)
        embedding_tokens_count = 0
        embedding = None
        if embedding_data is not None:
            if isinstance(embedding_data, str):
                embedding = self.load_embedding(embedding_data)
                if embedding is not None:
                    embedding_tokens_count = embedding.shape[0]
                    embedding = np.expand_dims(embedding, axis=0)
                else:
                    raise ValueError(
                        f"failed to load embedding file: {embedding_data}.")
        pad_token_id = 49407
        context = get_weighted_text_embeddings(self.tokenizer, self.text_clip_embedding, self.text_encoder, prompt,
                                               model_max_length=MAX_PROMPT_LENGTH,
                                               embedding=embedding, embedding_tokens_count=embedding_tokens_count,
                                               pad_token_id=pad_token_id)
        return context

    def gaussian_blur(self, image, radius=3, h_axis=1, v_axis=2):
        def build_filter1d(kernel_size):
            if kernel_size == 1:
                filter1d = [1]
            else:
                triangle = [[1, 1]]
                for i in range(1, kernel_size - 1):
                    cur_row = [1]
                    prev_row = triangle[i - 1]
                    for j in range(len(prev_row) - 1):
                        cur_row.append(prev_row[j] + prev_row[j + 1])
                    cur_row.append(1)
                    triangle.append(cur_row)
                filter1d = triangle[-1]
            filter1d = np.reshape(filter1d, (kernel_size,))
            return filter1d / np.sum(filter1d)

        weights = build_filter1d(radius)
        # Apply filter horizontally
        blurred_image = correlate1d(image, weights, axis=h_axis, output=None, mode="reflect", cval=0.0, origin=0)
        # Apply filter vertically
        blurred_image = correlate1d(blurred_image, weights, axis=v_axis, output=None, mode="reflect", cval=0.0,
                                    origin=0)
        return blurred_image

    @staticmethod
    def resize(image_array, new_h=None, new_w=None, mode="L"):
        if mode == "L":
            image_array = np.expand_dims(image_array, axis=-1)
        h, w, c = image_array.shape
        if new_h == h and new_w == w:
            return image_array
        h_bounds = 0, h - 1
        w_bounds = 0, w - 1
        y = np.expand_dims(np.linspace(h_bounds[0], h_bounds[1], new_h), axis=-1)
        x = np.expand_dims(np.linspace(w_bounds[0], w_bounds[1], new_w), axis=0)
        # Calculate the floor and ceiling values of x and y
        x_floor = np.floor(x).astype(int)
        x_ceil = np.ceil(x).astype(int)
        y_floor = np.floor(y).astype(int)
        y_ceil = np.ceil(y).astype(int)
        # Clip the values to stay within the image bounds
        x_floor = np.clip(x_floor, w_bounds[0], w_bounds[1])
        x_ceil = np.clip(x_ceil, w_bounds[0], w_bounds[1])
        y_floor = np.clip(y_floor, h_bounds[0], h_bounds[1])
        y_ceil = np.clip(y_ceil, h_bounds[0], h_bounds[1])
        # Calculate the fractional part of x and y
        dx = x - x_floor
        dy = y - y_floor
        # Get the values of the four neighboring pixels
        dx = np.expand_dims(dx, axis=-1)
        dy = np.expand_dims(dy, axis=-1)
        q11 = image_array[y_floor, x_floor, :]
        q21 = image_array[y_floor, x_ceil, :]
        q12 = image_array[y_ceil, x_floor, :]
        q22 = image_array[y_ceil, x_ceil, :]
        # Perform bilinear interpolation
        top_interp = q11 * (1.0 - dx) + q21 * dx
        bottom_interp = q12 * (1.0 - dx) + q22 * dx
        interpolated = top_interp * (1.0 - dy) + bottom_interp * dy
        if mode == "L":
            interpolated = np.squeeze(interpolated, axis=-1)
        return interpolated

    def preprocessed_image(self, x):
        if type(x) is str:
            image_array = Image.open(x).convert("RGB")
            image_array = image_array.resize((self.img_width, self.img_height))
        elif type(x) is np.ndarray:
            image_array = self.resize(x, self.img_height, self.img_width, mode="RGB")
        else:
            return None, None
        x = np.array(image_array, dtype=np.float32) / 255.0
        input_image_array = x[None, ..., :3]
        input_image_tensor = input_image_array * 2.0 - 1.0
        return input_image_array, input_image_tensor

    def preprocessed_mask(self, x, blur_radius=5):
        if type(x) is str:
            mask_array = Image.open(x).convert("L")
            mask_array = mask_array.resize((self.img_width, self.img_height))
        elif type(x) is np.ndarray:
            mask_array = self.resize(x, self.img_height, self.img_width, mode="L")
        else:
            return None, None
        input_mask_array = np.array(mask_array, dtype=np.float32) / 255.0
        input_mask_array = self.gaussian_blur(input_mask_array, radius=blur_radius, h_axis=0, v_axis=1)
        latent_mask_tensor = self.resize(input_mask_array, self.img_width // 8, self.img_height // 8,
                                         mode="L")
        return input_mask_array[None, ..., None], latent_mask_tensor[None, ..., None]

    def generate_image(
            self,
            encoded_text,
            negative_prompt=None,
            batch_size=1,
            num_steps=50,
            unconditional_guidance_scale=7.5,
            diffusion_noise=None,
            seed=None,
            negative_embedding=None,
            control_net_image=None,
            inpaint_mask=None,
            mask_blur_strength=5,
            reference_image=None,
            reference_image_strength=0.8,
    ):
        """Generates an image based on encoded text.

        The encoding passed to this method should be derived from
        `StableDiffusion.encode_text`.

        Args:
            encoded_text: Tensor of shape (`batch_size`, 77, 768), or a Tensor
                of shape (77, 768). When the batch axis is omitted, the same
                encoded text will be used to produce every generated image.
            batch_size: int, number of images to generate, defaults to 1.
            negative_prompt: a string containing information to negatively guide
                the image generation (e.g. by removing or altering certain
                aspects of the generated image), defaults to None.
            num_steps: int, number of diffusion steps (controls image quality),
                defaults to 50.
            unconditional_guidance_scale: float, controlling how closely the
                image should adhere to the prompt. Larger values result in more
                closely adhering to the prompt, but will make the image noisier.
                Defaults to 7.5.
            diffusion_noise: Tensor of shape (`batch_size`, img_height // 8,
                img_width // 8, 4), or a Tensor of shape (img_height // 8,
                img_width // 8, 4). Optional custom noise to seed the diffusion
                process. When the batch axis is omitted, the same noise will be
                used to seed diffusion for every generated image.
            seed: integer which is used to seed the random generation of
                diffusion noise, only to be specified if `diffusion_noise` is
                None.

        Example:

        ```python
        from keras_cv.models import StableDiffusion

        batch_size = 8
        model = StableDiffusion(img_height=512, img_width=512, jit_compile=True)
        e_tacos = model.encode_text("Tacos at dawn")
        e_watermelons = model.encode_text("Watermelons at dusk")

        e_interpolated = tf.linspace(e_tacos, e_watermelons, batch_size)
        images = model.generate_image(e_interpolated, batch_size=batch_size)
        ```
        """
        if diffusion_noise is not None and seed is not None:
            raise ValueError(
                "`diffusion_noise` and `seed` should not both be passed to "
                "`generate_image`. `seed` is only used to generate diffusion "
                "noise when it's not already user-specified."
            )

        context = self._expand_tensor(encoded_text, batch_size)
        if negative_prompt is None and negative_embedding is None:
            unconditional_context = np.repeat(
                self._get_unconditional_context(), batch_size, axis=0)
        else:
            if negative_prompt is None:
                negative_prompt = ""
            unconditional_context = self.encode_text(negative_prompt, negative_embedding)
            unconditional_context = self._expand_tensor(unconditional_context, batch_size)

        if diffusion_noise is not None:
            diffusion_noise = np.squeeze(diffusion_noise)
            if len(diffusion_noise.shape) == 3:
                diffusion_noise = np.repeat(np.expand_dims(diffusion_noise, axis=0), batch_size, axis=0)
        # Iterative reverse diffusion stage
        timesteps = np.linspace(0, 1000 - 1, num_steps, dtype=np.int32)
        init_time = None
        init_latent = None
        input_image_array = None
        input_mask_array = None
        latent_mask_tensor = None
        if inpaint_mask is not None:
            input_mask_array, latent_mask_tensor = self.preprocessed_mask(inpaint_mask, mask_blur_strength)
            if input_mask_array is None or latent_mask_tensor is None:
                print("wrong inpaint mask:{}".format(inpaint_mask))
        if reference_image is not None and (0. < reference_image_strength < 1.):
            input_image_array, input_image_tensor = self.preprocessed_image(reference_image)
            if input_image_tensor is not None:
                num_steps = int(num_steps * reference_image_strength + 0.5)
                init_time = timesteps[num_steps]
                init_latent = self.image_encoder.predict_on_batch(input_image_tensor)
                timesteps = timesteps[:num_steps]
            else:
                print("wrong reference image:{}".format(reference_image))
        latent = self._get_initial_diffusion_latent(batch_size=batch_size,
                                                    init_latent=init_latent,
                                                    init_time=init_time,
                                                    seed=seed,
                                                    noise=diffusion_noise)
        signal_rates, noise_rates, next_signal_rates, next_noise_rates = self._get_initial_alphas(timesteps)
        progbar = tf.keras.utils.Progbar(len(timesteps))
        iteration = 0
        hint_img = None
        if control_net_image is not None:
            if type(control_net_image) is str:
                image_array = Image.open(control_net_image).convert("RGB")
                image_array = image_array.resize((self.img_width, self.img_height))
            elif type(control_net_image) is np.ndarray:
                image_array = self.resize(control_net_image, self.img_height, self.img_width, mode="RGB")
            else:
                print("wrong controlnet image:{}".format(control_net_image))
                image_array = None
            if image_array is not None:
                controlnet_image = np.expand_dims(np.array(image_array, dtype=np.float32) / 255.0, axis=0)
                controlnet_image = np.tile(controlnet_image, (batch_size, 1, 1, 1))
                hint_img = self.hint_net.predict_on_batch(controlnet_image)
            else:
                control_net_image = None
        for index, timestep in list(enumerate(timesteps))[::-1]:
            latent_prev = latent  # Set aside the previous latent vector
            t_emb = self._get_timestep_embedding(timestep, batch_size)
            if control_net_image is not None and hint_img is not None:
                unconditional_controls = self.control_net.predict_on_batch(
                    [latent, t_emb, unconditional_context, hint_img])
                unconditional_latent = self.diffusion_model.predict_on_batch(
                    [latent, t_emb, unconditional_context] + list(unconditional_controls))
                controls = self.control_net.predict_on_batch([latent, t_emb, context, hint_img])
                latent = self.diffusion_model.predict_on_batch([latent, t_emb, context] + list(controls))
            else:
                unconditional_latent = self.diffusion_model.predict_on_batch(
                    [latent, t_emb, unconditional_context])
                latent = self.diffusion_model.predict_on_batch(
                    [latent, t_emb, context])
            latent = unconditional_latent + unconditional_guidance_scale * (
                    latent - unconditional_latent)
            pred_x0 = (latent_prev - noise_rates[index] * latent) / signal_rates[index]
            latent = (latent * next_noise_rates[index] + next_signal_rates[index] * pred_x0)
            if latent_mask_tensor is not None and init_latent is not None:
                latent_orgin = self._get_initial_diffusion_latent(batch_size=batch_size,
                                                                  init_latent=init_latent,
                                                                  init_time=timestep,
                                                                  seed=seed,
                                                                  noise=diffusion_noise)
                latent = latent_orgin * (1. - latent_mask_tensor) + latent * latent_mask_tensor
            iteration += 1
            progbar.update(iteration)

        # Decoding stage
        decoded = self.image_decoder.predict_on_batch(latent)
        decoded = np.array(((decoded + 1.) * 0.5), dtype=np.float32)
        if input_mask_array is not None and input_image_array is not None:
            decoded = input_image_array * (1. - input_mask_array) + decoded * input_mask_array
        return np.clip(decoded * 255., 0, 255).astype("uint8")

    def _get_unconditional_context(self):
        unconditional_tokens_ids = [49406] + [49407] * (MAX_PROMPT_LENGTH - 1)
        unconditional_tokens = np.asarray([unconditional_tokens_ids], dtype=np.int32)
        clip_embedding = self.text_clip_embedding.predict_on_batch([unconditional_tokens, self._get_pos_ids()])
        unconditional_context = self.text_encoder.predict_on_batch(clip_embedding)
        return unconditional_context

    def _expand_tensor(self, text_embedding, batch_size):
        """Extends a tensor by repeating it to fit the shape of the given batch
        size."""
        text_embedding = np.squeeze(text_embedding)
        if len(text_embedding.shape) == 2:
            text_embedding = np.repeat(
                np.expand_dims(text_embedding, axis=0), batch_size, axis=0
            )
        return text_embedding

    @property
    def image_encoder(self):
        """image_encoder returns the VAE Encoder with pretrained weights.

        Usage:
        ```python
        sd = keras_cv.models.StableDiffusion()
        my_image = np.ones((512, 512, 3))
        latent_representation = sd.image_encoder.predict(my_image)
        ```
        """
        if self._image_encoder is None:
            self._image_encoder = ImageEncoder()
            if self.jit_compile:
                self._image_encoder.compile(jit_compile=True)
        return self._image_encoder

    @property
    def hint_net(self):
        pass

    @property
    def control_net(self):
        pass

    @property
    def text_clip_embedding(self):
        pass

    @property
    def text_encoder(self):
        pass

    @property
    def diffusion_model(self):
        pass

    @property
    def image_decoder(self):
        """decoder returns the diffusion image decoder model with pretrained
        weights. Can be overriden for tasks where the decoder needs to be
        modified.
        """
        if self._image_decoder is None:
            self._image_decoder = ImageDecoder(self.img_height, self.img_width)
            if self.jit_compile:
                self._image_decoder.compile(jit_compile=True)
        return self._image_decoder

    @property
    def tokenizer(self):
        """tokenizer returns the tokenizer used for text inputs.
        Can be overriden for tasks like textual inversion where the tokenizer
        needs to be modified.
        """
        if self._tokenizer is None:
            self._tokenizer = SimpleTokenizer()
        return self._tokenizer

    def _get_timestep_embedding(
            self, timestep, batch_size, dim=320, max_period=10000
    ):
        half = dim // 2
        freqs = np.exp(
            -np.log(max_period) * np.asarray(range(0, half), dtype=np.float32) / half
        )
        args = np.asarray([timestep], dtype=np.float32) * freqs
        embedding = np.concatenate([np.cos(args), np.sin(args)], axis=0)
        embedding = np.reshape(embedding, [1, -1])
        return np.repeat(embedding, batch_size, axis=0)

    def _get_initial_alphas(self, timesteps):
        signal_rates = [self.signal_rates[t] for t in timesteps]
        noise_rates = [self.noise_rates[t] for t in timesteps]
        next_signal_rates = [1.0] + signal_rates[:-1]
        next_noise_rates = [0.0] + noise_rates[:-1]
        return signal_rates, noise_rates, next_signal_rates, next_noise_rates

    def _get_initial_diffusion_noise(self, batch_size, seed):
        if seed is not None:
            return tf.random.stateless_normal(
                (batch_size, self.img_height // 8, self.img_width // 8, 4),
                seed=[seed, seed],
            )
        else:
            return tf.random.normal(
                (batch_size, self.img_height // 8, self.img_width // 8, 4)
            )

    def _get_initial_diffusion_latent(self, batch_size, init_latent=None, init_time=None, seed=None,
                                      noise=None):
        if noise is None:
            noise = self._get_initial_diffusion_noise(batch_size, seed=seed)
        if init_latent is None:
            latent = noise
        else:
            latent = self.signal_rates[init_time] * np.repeat(init_latent, batch_size, axis=0) + self.noise_rates[
                init_time] * noise
        return latent

    @staticmethod
    def _get_pos_ids():
        return np.asarray([list(range(MAX_PROMPT_LENGTH))], dtype=np.int32)


class StableDiffusion(StableDiffusionBase):
    """Keras implementation of Stable Diffusion.

    Note that the StableDiffusion API, as well as the APIs of the sub-components
    of StableDiffusion (e.g. ImageEncoder, DiffusionModel) should be considered
    unstable at this point. We do not guarantee backwards compatability for
    future changes to these APIs.

    Stable Diffusion is a powerful image generation model that can be used,
    among other things, to generate pictures according to a short text
    description (called a "prompt").

    Arguments:
        img_height: int, height of the images to generate, in pixel. Note that
            only multiples of 128 are supported; the value provided will be
            rounded to the nearest valid value. Defaults to 512.
        img_width: int, width of the images to generate, in pixel. Note that
            only multiples of 128 are supported; the value provided will be
            rounded to the nearest valid value. Defaults to 512.
        jit_compile: bool, whether to compile the underlying models to XLA.
            This can lead to a significant speedup on some systems. Defaults to
            False.

    Example:

    ```python
    from keras_cv.models import StableDiffusion
    from PIL import Image

    model = StableDiffusion(img_height=512, img_width=512, jit_compile=True)
    img = model.text_to_image(
        prompt="A beautiful horse running through a field",
        batch_size=1,  # How many images to generate at once
        num_steps=25,  # Number of iterations (controls image quality)
        seed=123,  # Set this to always get the same image from the same prompt
    )
    Image.fromarray(img[0]).save("horse.png")
    print("saved at horse.png")
    ```

    References:
    - [About Stable Diffusion](https://stability.ai/blog/stable-diffusion-announcement)
    - [Original implementation](https://github.com/CompVis/stable-diffusion)
    """  # noqa: E501

    def __init__(
            self,
            img_height=512,
            img_width=512,
            jit_compile=False,
            clip_skip=-1,
            civitai_model=None,
            lora_path=None,
            controlnet_path=None):
        super().__init__(img_height, img_width, jit_compile)
        self.clip_skip = clip_skip
        self.civitai_model = civitai_model
        self.controlnet_path = controlnet_path
        self.lora_path = None
        self.text_encoder_lora_dict = None
        self.unet_lora_dict = None
        if os.path.exists(str(lora_path)):
            self.text_encoder_lora_dict, self.unet_lora_dict = load_weights_from_lora(lora_path)
            self.lora_path = lora_path
        print(
            "By using this model checkpoint, you acknowledge that its usage is "
            "subject to the terms of the CreativeML Open RAIL-M license at "
            "https://raw.githubusercontent.com/CompVis/stable-diffusion/main/LICENSE"  # noqa: E501
        )

    @property
    def text_encoder(self):
        """text_encoder returns the text encoder with pretrained weights.
        Can be overriden for tasks like textual inversion where the text encoder
        needs to be modified.
        """
        if self._text_encoder is None:
            self._text_encoder = TextEncoder(MAX_PROMPT_LENGTH, clip_skip=self.clip_skip,
                                             lora_dict=self.text_encoder_lora_dict)
            if self.jit_compile:
                self._text_encoder.compile(jit_compile=True)
        return self._text_encoder

    @property
    def hint_net(self):
        if self._hint_net is None:
            self._hint_net = HintNet(self.img_height, self.img_width, controlnet_path=self.controlnet_path)
            if self.jit_compile:
                self._hint_net.compile(jit_compile=True)
        return self._hint_net

    @property
    def control_net(self):
        if self._control_net is None:
            self._control_net = ControlNet(self.img_height, self.img_width, controlnet_path=self.controlnet_path)
            if self.jit_compile:
                self._control_net.compile(jit_compile=True)
        return self._control_net

    @property
    def text_clip_embedding(self):
        """text_clip_embedding returns the text encoder with pretrained weights.
        Can be overriden for tasks like textual inversion where the text encoder
        needs to be modified.
        """
        if self._text_clip_embedding is None:
            self._text_clip_embedding = TextClipEmbedding(MAX_PROMPT_LENGTH)
            if self.jit_compile:
                self._text_clip_embedding.compile(jit_compile=True)
        return self._text_clip_embedding

    @property
    def diffusion_model(self):
        """diffusion_model returns the diffusion model with pretrained weights.
        Can be overriden for tasks where the diffusion model needs to be
        modified.
        """
        if self._diffusion_model is None:
            self._diffusion_model = DiffusionModel(
                self.img_height, self.img_width, ckpt_path=self.civitai_model,
                apply_control_net=self.controlnet_path is not None, lora_dict=self.unet_lora_dict)
            if self.jit_compile:
                self._diffusion_model.compile(jit_compile=True)
        return self._diffusion_model

stable_diffusion/stable_diffusion.py
cpuimage/minSDTF
# Copyright 2022 The KerasCV Authors
#
# Licensed under the Apache License, Version 2.0 (the "License");
# you may not use this file except in compliance with the License.
# You may obtain a copy of the License at
#
#     https://www.apache.org/licenses/LICENSE-2.0
#
# Unless required by applicable law or agreed to in writing, software
# distributed under the License is distributed on an "AS IS" BASIS,
# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
# See the License for the specific language governing permissions and
# limitations under the License.
import os

import numpy as np
import tensorflow as tf

from .ckpt_loader import load_weights_from_file


class CLIPEmbedding(tf.keras.layers.Layer):
    def __init__(self, input_dim=49408, output_dim=768, max_length=77, **kwargs):
        super().__init__(**kwargs)
        self.token_embedding = tf.keras.layers.Embedding(input_dim, output_dim, name="token_embedding")
        self.position_embedding = tf.keras.layers.Embedding(max_length, output_dim, name="position_embedding")

    def call(self, inputs):
        tokens, positions = inputs
        tokens = self.token_embedding(tokens)
        positions = self.position_embedding(positions)
        return tokens + positions


class CLIPEncoderLayer(tf.keras.layers.Layer):
    def __init__(self, embed_dim, num_heads, activation=None, **kwargs):
        super().__init__(**kwargs)
        self.layer_norm1 = tf.keras.layers.LayerNormalization(epsilon=1e-5, name="layer_norm1")
        self.clip_attn = CLIPAttention(embed_dim, num_heads, causal=True, name="self_attn")
        self.layer_norm2 = tf.keras.layers.LayerNormalization(epsilon=1e-5, name="layer_norm2")
        self.fc1 = tf.keras.layers.Dense(embed_dim * 4, name="mlp.fc1")
        self.fc2 = tf.keras.layers.Dense(embed_dim, name="mlp.fc2")
        self.activation = activation

    def call(self, inputs):
        residual = inputs
        x = self.layer_norm1(inputs)
        x = self.clip_attn(x)
        x = residual + x
        residual = x
        x = self.layer_norm2(x)
        x = self.fc1(x)
        x = self.activation(x)
        x = self.fc2(x)
        return x + residual


class CLIPAttention(tf.keras.layers.Layer):
    def __init__(self, embed_dim=768, num_heads=12, causal=True, **kwargs):
        super().__init__(**kwargs)
        self.embed_dim = embed_dim
        self.num_heads = num_heads
        self.causal = causal
        self.head_dim = self.embed_dim // self.num_heads
        self.scale = self.head_dim ** -0.5
        self.q_proj = tf.keras.layers.Dense(self.embed_dim, name="q_proj")
        self.k_proj = tf.keras.layers.Dense(self.embed_dim, name="k_proj")
        self.v_proj = tf.keras.layers.Dense(self.embed_dim, name="v_proj")
        self.out_proj = tf.keras.layers.Dense(self.embed_dim, name="out_proj")

    def reshape_states(self, x, sequence_length, batch_size):
        x = tf.reshape(
            x, (batch_size, sequence_length, self.num_heads, self.head_dim))
        return tf.transpose(x, (0, 2, 1, 3))  # bs, heads, sequence_length, head_dim

    def call(self, inputs, attention_mask=None):
        if attention_mask is None and self.causal:
            length = inputs.get_shape().as_list()[1]
            attention_mask = tf.cast(np.triu(np.ones((1, 1, length, length), dtype="float32") * -np.inf, k=1),
                                     dtype=self.compute_dtype)
        _, tgt_len, embed_dim = inputs.shape
        query_states = self.q_proj(inputs) * self.scale
        key_states = self.reshape_states(self.k_proj(inputs), tgt_len, -1)
        value_states = self.reshape_states(self.v_proj(inputs), tgt_len, -1)
        proj_shape = (-1, tgt_len, self.head_dim)
        query_states = self.reshape_states(query_states, tgt_len, -1)
        query_states = tf.reshape(query_states, proj_shape)
        key_states = tf.reshape(key_states, proj_shape)
        src_len = tgt_len
        value_states = tf.reshape(value_states, proj_shape)
        attn_weights = query_states @ tf.transpose(key_states, (0, 2, 1))
        attn_weights = tf.reshape(attn_weights, (-1, self.num_heads, tgt_len, src_len))
        attn_weights = attn_weights + attention_mask
        attn_weights = tf.reshape(attn_weights, (-1, tgt_len, src_len))
        attn_weights = tf.nn.softmax(attn_weights)
        attn_output = attn_weights @ value_states
        attn_output = tf.reshape(attn_output, (-1, self.num_heads, tgt_len, self.head_dim))
        attn_output = tf.transpose(attn_output, (0, 2, 1, 3))
        attn_output = tf.reshape(attn_output, (-1, tgt_len, embed_dim))
        return self.out_proj(attn_output)


def quick_gelu(x):
    return x * tf.sigmoid(x * 1.702)


class TextClipEmbedding(tf.keras.Model):
    def __init__(self, max_length, embed_dim=768, vocab_size=49408, name=None, ckpt_path=None):
        tokens = tf.keras.layers.Input(shape=(max_length,), dtype="int32", name="tokens")
        positions = tf.keras.layers.Input(shape=(max_length,), dtype="int32", name="positions")
        clip_emb = CLIPEmbedding(vocab_size, embed_dim, max_length, name="embeddings")([tokens, positions])
        super().__init__([tokens, positions], clip_emb, name=name)
        origin = "https://huggingface.co/runwayml/stable-diffusion-v1-5/resolve/main/text_encoder/model.safetensors"
        ckpt_mapping = [('text_model.embeddings.token_embedding.weight', None),
                        ('text_model.embeddings.position_embedding.weight', None)]
        if ckpt_path is not None:
            if os.path.exists(ckpt_path):
                load_weights_from_file(self, ckpt_path, ckpt_mapping=ckpt_mapping)
                return
            else:
                origin = ckpt_path
        model_weights_fpath = tf.keras.utils.get_file(origin=origin)
        if os.path.exists(model_weights_fpath):
            load_weights_from_file(self, model_weights_fpath, ckpt_mapping=ckpt_mapping)


class TextEncoder(tf.keras.Model):
    def __init__(self, max_length=77, embed_dim=768, num_heads=12, num_layers=12, clip_skip=-2, name=None,
                 ckpt_path=None, lora_dict=None):
        clip_emb = tf.keras.layers.Input(shape=(max_length, embed_dim), dtype="float32", name="clip_emb")
        x = clip_emb
        out = []
        for idx in range(num_layers):
            x = CLIPEncoderLayer(embed_dim, num_heads, activation=quick_gelu,
                                 name="text_model.encoder.layers.{}".format(idx))(x)
            out.append(x)
        embedded = tf.keras.layers.LayerNormalization(epsilon=1e-5, name="text_model.final_layer_norm")(out[clip_skip])
        super().__init__(clip_emb, embedded, name=name)
        origin = "https://huggingface.co/runwayml/stable-diffusion-v1-5/resolve/main/text_encoder/model.safetensors"
        ckpt_mapping = []
        for idx in range(0, num_layers + clip_skip + 1):
            layers_name = 'text_model.encoder.layers.{}'.format(idx)
            ckpt_mapping.append(('{}.layer_norm1.weight'.format(layers_name), None))
            ckpt_mapping.append(('{}.layer_norm1.bias'.format(layers_name), None))
            ckpt_mapping.append(('{}.self_attn.q_proj.weight'.format(layers_name), (1, 0)))
            ckpt_mapping.append(('{}.self_attn.q_proj.bias'.format(layers_name), None))
            ckpt_mapping.append(('{}.self_attn.k_proj.weight'.format(layers_name), (1, 0)))
            ckpt_mapping.append(('{}.self_attn.k_proj.bias'.format(layers_name), None))
            ckpt_mapping.append(('{}.self_attn.v_proj.weight'.format(layers_name), (1, 0)))
            ckpt_mapping.append(('{}.self_attn.v_proj.bias'.format(layers_name), None))
            ckpt_mapping.append(('{}.self_attn.out_proj.weight'.format(layers_name), (1, 0)))
            ckpt_mapping.append(('{}.self_attn.out_proj.bias'.format(layers_name), None))
            ckpt_mapping.append(('{}.layer_norm2.weight'.format(layers_name), None))
            ckpt_mapping.append(('{}.layer_norm2.bias'.format(layers_name), None))
            ckpt_mapping.append(('{}.mlp.fc1.weight'.format(layers_name), (1, 0)))
            ckpt_mapping.append(('{}.mlp.fc1.bias'.format(layers_name), None))
            ckpt_mapping.append(('{}.mlp.fc2.weight'.format(layers_name), (1, 0)))
            ckpt_mapping.append(('{}.mlp.fc2.bias'.format(layers_name), None))
        ckpt_mapping.append(('text_model.final_layer_norm.weight', None))
        ckpt_mapping.append(('text_model.final_layer_norm.bias', None))
        if ckpt_path is not None:
            if os.path.exists(ckpt_path):
                load_weights_from_file(self, ckpt_path, ckpt_mapping=ckpt_mapping, lora_dict=lora_dict)
                return
            else:
                origin = ckpt_path
        model_weights_fpath = tf.keras.utils.get_file(origin=origin)
        if os.path.exists(model_weights_fpath):
            load_weights_from_file(self, model_weights_fpath, ckpt_mapping=ckpt_mapping, lora_dict=lora_dict)

stable_diffusion/text_encoder.py
cpuimage/minSDTF
from PIL import Image

from stable_diffusion.stable_diffusion import StableDiffusion

model = StableDiffusion(img_height=512, img_width=512, jit_compile=True)
img = model.text_to_image(
    "a cute girl.",
    num_steps=25,
    seed=123456,
)
Image.fromarray(img[0]).save("girl.jpg")
print("Saved at girl.jpg")

text_to_image.py
cpuimage/minSDTF