# coding=utf-8 # Copyright 2024 The HuggingFace Inc. team. # # 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 # # http://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. """ Conversion script for the T2I-Adapter checkpoints. """ import argparse import torch from diffusers import T2IAdapter def convert_adapter(src_state, in_channels): original_body_length = max([int(x.split(".")[1]) for x in src_state.keys() if "body." in x]) + 1 assert original_body_length == 8 # (0, 1) -> channels 1 assert src_state["body.0.block1.weight"].shape == (320, 320, 3, 3) # (2, 3) -> channels 2 assert src_state["body.2.in_conv.weight"].shape == (640, 320, 1, 1) # (4, 5) -> channels 3 assert src_state["body.4.in_conv.weight"].shape == (1280, 640, 1, 1) # (6, 7) -> channels 4 assert src_state["body.6.block1.weight"].shape == (1280, 1280, 3, 3) res_state = { "adapter.conv_in.weight": src_state.pop("conv_in.weight"), "adapter.conv_in.bias": src_state.pop("conv_in.bias"), # 0.resnets.0 "adapter.body.0.resnets.0.block1.weight": src_state.pop("body.0.block1.weight"), "adapter.body.0.resnets.0.block1.bias": src_state.pop("body.0.block1.bias"), "adapter.body.0.resnets.0.block2.weight": src_state.pop("body.0.block2.weight"), "adapter.body.0.resnets.0.block2.bias": src_state.pop("body.0.block2.bias"), # 0.resnets.1 "adapter.body.0.resnets.1.block1.weight": src_state.pop("body.1.block1.weight"), "adapter.body.0.resnets.1.block1.bias": src_state.pop("body.1.block1.bias"), "adapter.body.0.resnets.1.block2.weight": src_state.pop("body.1.block2.weight"), "adapter.body.0.resnets.1.block2.bias": src_state.pop("body.1.block2.bias"), # 1 "adapter.body.1.in_conv.weight": src_state.pop("body.2.in_conv.weight"), "adapter.body.1.in_conv.bias": src_state.pop("body.2.in_conv.bias"), # 1.resnets.0 "adapter.body.1.resnets.0.block1.weight": src_state.pop("body.2.block1.weight"), "adapter.body.1.resnets.0.block1.bias": src_state.pop("body.2.block1.bias"), "adapter.body.1.resnets.0.block2.weight": src_state.pop("body.2.block2.weight"), "adapter.body.1.resnets.0.block2.bias": src_state.pop("body.2.block2.bias"), # 1.resnets.1 "adapter.body.1.resnets.1.block1.weight": src_state.pop("body.3.block1.weight"), "adapter.body.1.resnets.1.block1.bias": src_state.pop("body.3.block1.bias"), "adapter.body.1.resnets.1.block2.weight": src_state.pop("body.3.block2.weight"), "adapter.body.1.resnets.1.block2.bias": src_state.pop("body.3.block2.bias"), # 2 "adapter.body.2.in_conv.weight": src_state.pop("body.4.in_conv.weight"), "adapter.body.2.in_conv.bias": src_state.pop("body.4.in_conv.bias"), # 2.resnets.0 "adapter.body.2.resnets.0.block1.weight": src_state.pop("body.4.block1.weight"), "adapter.body.2.resnets.0.block1.bias": src_state.pop("body.4.block1.bias"), "adapter.body.2.resnets.0.block2.weight": src_state.pop("body.4.block2.weight"), "adapter.body.2.resnets.0.block2.bias": src_state.pop("body.4.block2.bias"), # 2.resnets.1 "adapter.body.2.resnets.1.block1.weight": src_state.pop("body.5.block1.weight"), "adapter.body.2.resnets.1.block1.bias": src_state.pop("body.5.block1.bias"), "adapter.body.2.resnets.1.block2.weight": src_state.pop("body.5.block2.weight"), "adapter.body.2.resnets.1.block2.bias": src_state.pop("body.5.block2.bias"), # 3.resnets.0 "adapter.body.3.resnets.0.block1.weight": src_state.pop("body.6.block1.weight"), "adapter.body.3.resnets.0.block1.bias": src_state.pop("body.6.block1.bias"), "adapter.body.3.resnets.0.block2.weight": src_state.pop("body.6.block2.weight"), "adapter.body.3.resnets.0.block2.bias": src_state.pop("body.6.block2.bias"), # 3.resnets.1 "adapter.body.3.resnets.1.block1.weight": src_state.pop("body.7.block1.weight"), "adapter.body.3.resnets.1.block1.bias": src_state.pop("body.7.block1.bias"), "adapter.body.3.resnets.1.block2.weight": src_state.pop("body.7.block2.weight"), "adapter.body.3.resnets.1.block2.bias": src_state.pop("body.7.block2.bias"), } assert len(src_state) == 0 adapter = T2IAdapter(in_channels=in_channels, adapter_type="full_adapter") adapter.load_state_dict(res_state) return adapter def convert_light_adapter(src_state): original_body_length = max([int(x.split(".")[1]) for x in src_state.keys() if "body." in x]) + 1 assert original_body_length == 4 res_state = { # body.0.in_conv "adapter.body.0.in_conv.weight": src_state.pop("body.0.in_conv.weight"), "adapter.body.0.in_conv.bias": src_state.pop("body.0.in_conv.bias"), # body.0.resnets.0 "adapter.body.0.resnets.0.block1.weight": src_state.pop("body.0.body.0.block1.weight"), "adapter.body.0.resnets.0.block1.bias": src_state.pop("body.0.body.0.block1.bias"), "adapter.body.0.resnets.0.block2.weight": src_state.pop("body.0.body.0.block2.weight"), "adapter.body.0.resnets.0.block2.bias": src_state.pop("body.0.body.0.block2.bias"), # body.0.resnets.1 "adapter.body.0.resnets.1.block1.weight": src_state.pop("body.0.body.1.block1.weight"), "adapter.body.0.resnets.1.block1.bias": src_state.pop("body.0.body.1.block1.bias"), "adapter.body.0.resnets.1.block2.weight": src_state.pop("body.0.body.1.block2.weight"), "adapter.body.0.resnets.1.block2.bias": src_state.pop("body.0.body.1.block2.bias"), # body.0.resnets.2 "adapter.body.0.resnets.2.block1.weight": src_state.pop("body.0.body.2.block1.weight"), "adapter.body.0.resnets.2.block1.bias": src_state.pop("body.0.body.2.block1.bias"), "adapter.body.0.resnets.2.block2.weight": src_state.pop("body.0.body.2.block2.weight"), "adapter.body.0.resnets.2.block2.bias": src_state.pop("body.0.body.2.block2.bias"), # body.0.resnets.3 "adapter.body.0.resnets.3.block1.weight": src_state.pop("body.0.body.3.block1.weight"), "adapter.body.0.resnets.3.block1.bias": src_state.pop("body.0.body.3.block1.bias"), "adapter.body.0.resnets.3.block2.weight": src_state.pop("body.0.body.3.block2.weight"), "adapter.body.0.resnets.3.block2.bias": src_state.pop("body.0.body.3.block2.bias"), # body.0.out_conv "adapter.body.0.out_conv.weight": src_state.pop("body.0.out_conv.weight"), "adapter.body.0.out_conv.bias": src_state.pop("body.0.out_conv.bias"), # body.1.in_conv "adapter.body.1.in_conv.weight": src_state.pop("body.1.in_conv.weight"), "adapter.body.1.in_conv.bias": src_state.pop("body.1.in_conv.bias"), # body.1.resnets.0 "adapter.body.1.resnets.0.block1.weight": src_state.pop("body.1.body.0.block1.weight"), "adapter.body.1.resnets.0.block1.bias": src_state.pop("body.1.body.0.block1.bias"), "adapter.body.1.resnets.0.block2.weight": src_state.pop("body.1.body.0.block2.weight"), "adapter.body.1.resnets.0.block2.bias": src_state.pop("body.1.body.0.block2.bias"), # body.1.resnets.1 "adapter.body.1.resnets.1.block1.weight": src_state.pop("body.1.body.1.block1.weight"), "adapter.body.1.resnets.1.block1.bias": src_state.pop("body.1.body.1.block1.bias"), "adapter.body.1.resnets.1.block2.weight": src_state.pop("body.1.body.1.block2.weight"), "adapter.body.1.resnets.1.block2.bias": src_state.pop("body.1.body.1.block2.bias"), # body.1.body.2 "adapter.body.1.resnets.2.block1.weight": src_state.pop("body.1.body.2.block1.weight"), "adapter.body.1.resnets.2.block1.bias": src_state.pop("body.1.body.2.block1.bias"), "adapter.body.1.resnets.2.block2.weight": src_state.pop("body.1.body.2.block2.weight"), "adapter.body.1.resnets.2.block2.bias": src_state.pop("body.1.body.2.block2.bias"), # body.1.body.3 "adapter.body.1.resnets.3.block1.weight": src_state.pop("body.1.body.3.block1.weight"), "adapter.body.1.resnets.3.block1.bias": src_state.pop("body.1.body.3.block1.bias"), "adapter.body.1.resnets.3.block2.weight": src_state.pop("body.1.body.3.block2.weight"), "adapter.body.1.resnets.3.block2.bias": src_state.pop("body.1.body.3.block2.bias"), # body.1.out_conv "adapter.body.1.out_conv.weight": src_state.pop("body.1.out_conv.weight"), "adapter.body.1.out_conv.bias": src_state.pop("body.1.out_conv.bias"), # body.2.in_conv "adapter.body.2.in_conv.weight": src_state.pop("body.2.in_conv.weight"), "adapter.body.2.in_conv.bias": src_state.pop("body.2.in_conv.bias"), # body.2.body.0 "adapter.body.2.resnets.0.block1.weight": src_state.pop("body.2.body.0.block1.weight"), "adapter.body.2.resnets.0.block1.bias": src_state.pop("body.2.body.0.block1.bias"), "adapter.body.2.resnets.0.block2.weight": src_state.pop("body.2.body.0.block2.weight"), "adapter.body.2.resnets.0.block2.bias": src_state.pop("body.2.body.0.block2.bias"), # body.2.body.1 "adapter.body.2.resnets.1.block1.weight": src_state.pop("body.2.body.1.block1.weight"), "adapter.body.2.resnets.1.block1.bias": src_state.pop("body.2.body.1.block1.bias"), "adapter.body.2.resnets.1.block2.weight": src_state.pop("body.2.body.1.block2.weight"), "adapter.body.2.resnets.1.block2.bias": src_state.pop("body.2.body.1.block2.bias"), # body.2.body.2 "adapter.body.2.resnets.2.block1.weight": src_state.pop("body.2.body.2.block1.weight"), "adapter.body.2.resnets.2.block1.bias": src_state.pop("body.2.body.2.block1.bias"), "adapter.body.2.resnets.2.block2.weight": src_state.pop("body.2.body.2.block2.weight"), "adapter.body.2.resnets.2.block2.bias": src_state.pop("body.2.body.2.block2.bias"), # body.2.body.3 "adapter.body.2.resnets.3.block1.weight": src_state.pop("body.2.body.3.block1.weight"), "adapter.body.2.resnets.3.block1.bias": src_state.pop("body.2.body.3.block1.bias"), "adapter.body.2.resnets.3.block2.weight": src_state.pop("body.2.body.3.block2.weight"), "adapter.body.2.resnets.3.block2.bias": src_state.pop("body.2.body.3.block2.bias"), # body.2.out_conv "adapter.body.2.out_conv.weight": src_state.pop("body.2.out_conv.weight"), "adapter.body.2.out_conv.bias": src_state.pop("body.2.out_conv.bias"), # body.3.in_conv "adapter.body.3.in_conv.weight": src_state.pop("body.3.in_conv.weight"), "adapter.body.3.in_conv.bias": src_state.pop("body.3.in_conv.bias"), # body.3.body.0 "adapter.body.3.resnets.0.block1.weight": src_state.pop("body.3.body.0.block1.weight"), "adapter.body.3.resnets.0.block1.bias": src_state.pop("body.3.body.0.block1.bias"), "adapter.body.3.resnets.0.block2.weight": src_state.pop("body.3.body.0.block2.weight"), "adapter.body.3.resnets.0.block2.bias": src_state.pop("body.3.body.0.block2.bias"), # body.3.body.1 "adapter.body.3.resnets.1.block1.weight": src_state.pop("body.3.body.1.block1.weight"), "adapter.body.3.resnets.1.block1.bias": src_state.pop("body.3.body.1.block1.bias"), "adapter.body.3.resnets.1.block2.weight": src_state.pop("body.3.body.1.block2.weight"), "adapter.body.3.resnets.1.block2.bias": src_state.pop("body.3.body.1.block2.bias"), # body.3.body.2 "adapter.body.3.resnets.2.block1.weight": src_state.pop("body.3.body.2.block1.weight"), "adapter.body.3.resnets.2.block1.bias": src_state.pop("body.3.body.2.block1.bias"), "adapter.body.3.resnets.2.block2.weight": src_state.pop("body.3.body.2.block2.weight"), "adapter.body.3.resnets.2.block2.bias": src_state.pop("body.3.body.2.block2.bias"), # body.3.body.3 "adapter.body.3.resnets.3.block1.weight": src_state.pop("body.3.body.3.block1.weight"), "adapter.body.3.resnets.3.block1.bias": src_state.pop("body.3.body.3.block1.bias"), "adapter.body.3.resnets.3.block2.weight": src_state.pop("body.3.body.3.block2.weight"), "adapter.body.3.resnets.3.block2.bias": src_state.pop("body.3.body.3.block2.bias"), # body.3.out_conv "adapter.body.3.out_conv.weight": src_state.pop("body.3.out_conv.weight"), "adapter.body.3.out_conv.bias": src_state.pop("body.3.out_conv.bias"), } assert len(src_state) == 0 adapter = T2IAdapter(in_channels=3, channels=[320, 640, 1280], num_res_blocks=4, adapter_type="light_adapter") adapter.load_state_dict(res_state) return adapter if __name__ == "__main__": parser = argparse.ArgumentParser() parser.add_argument( "--checkpoint_path", default=None, type=str, required=True, help="Path to the checkpoint to convert." ) parser.add_argument( "--output_path", default=None, type=str, required=True, help="Path to the store the result checkpoint." ) parser.add_argument( "--is_adapter_light", action="store_true", help="Is checkpoint come from Adapter-Light architecture. ex: color-adapter", ) parser.add_argument("--in_channels", required=False, type=int, help="Input channels for non-light adapter") args = parser.parse_args() src_state = torch.load(args.checkpoint_path) if args.is_adapter_light: adapter = convert_light_adapter(src_state) else: if args.in_channels is None: raise ValueError("set `--in_channels=`") adapter = convert_adapter(src_state, args.in_channels) adapter.save_pretrained(args.output_path)