FilesExpand file tree

main

/

change_naming_configs_and_checkpoints.py

Copy path

Blame

More file actions

Blame

More file actions

Latest commit

 

History

History

History

112 lines (92 loc) · 3.81 KB

main

/

change_naming_configs_and_checkpoints.py

Copy path

Top

File metadata and controls

• Code
• Blame

112 lines (92 loc) · 3.81 KB

Raw

Copy raw file

Download raw file

Open symbols panel

Edit and raw actions

1

2

3

4

5

6

7

8

9

10

11

12

13

14

15

16

17

18

19

20

21

22

23

24

25

26

27

28

29

30

31

32

33

34

35

36

37

38

39

40

41

42

43

44

45

46

47

48

49

50

51

52

53

54

55

56

57

58

59

60

61

62

63

64

65

66

67

68

69

70

71

72

73

74

75

76

77

78

79

80

81

82

83

84

85

86

87

88

89

90

91

92

93

94

95

96

97

98

99

100

101

102

103

104

105

106

107

108

109

110

111

112

# coding=utf-8

# Copyright 2022 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 LDM checkpoints. """

import argparse

import os

import json

import torch

from diffusers import UNet2DModel, UNet2DConditionModel

from transformers.file_utils import has_file

do_only_config = False

do_only_weights = True

do_only_renaming = False

if __name__ == "__main__":

parser = argparse.ArgumentParser()

parser.add_argument(

"--repo_path",

default=None,

type=str,

required=True,

help="The config json file corresponding to the architecture.",

)

parser.add_argument(

"--dump_path", default=None, type=str, required=True, help="Path to the output model."

)

args = parser.parse_args()

config_parameters_to_change = {

"image_size": "sample_size",

"num_res_blocks": "layers_per_block",

"block_channels": "block_out_channels",

"down_blocks": "down_block_types",

"up_blocks": "up_block_types",

"downscale_freq_shift": "freq_shift",

"resnet_num_groups": "norm_num_groups",

"resnet_act_fn": "act_fn",

"resnet_eps": "norm_eps",

"num_head_channels": "attention_head_dim",

}

key_parameters_to_change = {

"time_steps": "time_proj",

"mid": "mid_block",

"downsample_blocks": "down_blocks",

"upsample_blocks": "up_blocks",

}

subfolder = "" if has_file(args.repo_path, "config.json") else "unet"

with open(os.path.join(args.repo_path, subfolder, "config.json"), "r", encoding="utf-8") as reader:

text = reader.read()

config = json.loads(text)

if do_only_config:

for key in config_parameters_to_change.keys():

config.pop(key, None)

if has_file(args.repo_path, "config.json"):

model = UNet2DModel(**config)

else:

class_name = UNet2DConditionModel if "ldm-text2im-large-256" in args.repo_path else UNet2DModel

model = class_name(**config)

if do_only_config:

model.save_config(os.path.join(args.repo_path, subfolder))

config = dict(model.config)

if do_only_renaming:

for key, value in config_parameters_to_change.items():

if key in config:

config[value] = config[key]

del config[key]

config["down_block_types"] = [k.replace("UNetRes", "") for k in config["down_block_types"]]

config["up_block_types"] = [k.replace("UNetRes", "") for k in config["up_block_types"]]

if do_only_weights:

state_dict = torch.load(os.path.join(args.repo_path, subfolder, "diffusion_pytorch_model.bin"))

new_state_dict = {}

for param_key, param_value in state_dict.items():

if param_key.endswith(".op.bias") or param_key.endswith(".op.weight"):

continue

has_changed = False

for key, new_key in key_parameters_to_change.items():

if not has_changed and param_key.split(".")[0] == key:

new_state_dict[".".join([new_key] + param_key.split(".")[1:])] = param_value

has_changed = True

if not has_changed:

new_state_dict[param_key] = param_value

model.load_state_dict(new_state_dict)

model.save_pretrained(os.path.join(args.repo_path, subfolder))