FilesExpand file tree

main

/

convert_vae_pt_to_diffusers.py

Copy path

Blame

More file actions

Blame

More file actions

Latest commit

 

History

History

History

151 lines (119 loc) · 6.98 KB

main

/

convert_vae_pt_to_diffusers.py

Copy path

Top

File metadata and controls

• Code
• Blame

151 lines (119 loc) · 6.98 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

113

114

115

116

117

118

119

120

121

122

123

124

125

126

127

128

129

130

131

132

133

134

135

136

137

138

139

140

141

142

143

144

145

146

147

148

149

150

151

import argparse

import io

import requests

import torch

from omegaconf import OmegaConf

from diffusers import AutoencoderKL

from diffusers.pipelines.stable_diffusion.convert_from_ckpt import (

assign_to_checkpoint,

conv_attn_to_linear,

create_vae_diffusers_config,

renew_vae_attention_paths,

renew_vae_resnet_paths,

)

def custom_convert_ldm_vae_checkpoint(checkpoint, config):

vae_state_dict = checkpoint

new_checkpoint = {}

new_checkpoint["encoder.conv_in.weight"] = vae_state_dict["encoder.conv_in.weight"]

new_checkpoint["encoder.conv_in.bias"] = vae_state_dict["encoder.conv_in.bias"]

new_checkpoint["encoder.conv_out.weight"] = vae_state_dict["encoder.conv_out.weight"]

new_checkpoint["encoder.conv_out.bias"] = vae_state_dict["encoder.conv_out.bias"]

new_checkpoint["encoder.conv_norm_out.weight"] = vae_state_dict["encoder.norm_out.weight"]

new_checkpoint["encoder.conv_norm_out.bias"] = vae_state_dict["encoder.norm_out.bias"]

new_checkpoint["decoder.conv_in.weight"] = vae_state_dict["decoder.conv_in.weight"]

new_checkpoint["decoder.conv_in.bias"] = vae_state_dict["decoder.conv_in.bias"]

new_checkpoint["decoder.conv_out.weight"] = vae_state_dict["decoder.conv_out.weight"]

new_checkpoint["decoder.conv_out.bias"] = vae_state_dict["decoder.conv_out.bias"]

new_checkpoint["decoder.conv_norm_out.weight"] = vae_state_dict["decoder.norm_out.weight"]

new_checkpoint["decoder.conv_norm_out.bias"] = vae_state_dict["decoder.norm_out.bias"]

new_checkpoint["quant_conv.weight"] = vae_state_dict["quant_conv.weight"]

new_checkpoint["quant_conv.bias"] = vae_state_dict["quant_conv.bias"]

new_checkpoint["post_quant_conv.weight"] = vae_state_dict["post_quant_conv.weight"]

new_checkpoint["post_quant_conv.bias"] = vae_state_dict["post_quant_conv.bias"]

# Retrieves the keys for the encoder down blocks only

num_down_blocks = len({".".join(layer.split(".")[:3]) for layer in vae_state_dict if "encoder.down" in layer})

down_blocks = {

layer_id: [key for key in vae_state_dict if f"down.{layer_id}" in key] for layer_id in range(num_down_blocks)

}

# Retrieves the keys for the decoder up blocks only

num_up_blocks = len({".".join(layer.split(".")[:3]) for layer in vae_state_dict if "decoder.up" in layer})

up_blocks = {

layer_id: [key for key in vae_state_dict if f"up.{layer_id}" in key] for layer_id in range(num_up_blocks)

}

for i in range(num_down_blocks):

resnets = [key for key in down_blocks[i] if f"down.{i}" in key and f"down.{i}.downsample" not in key]

if f"encoder.down.{i}.downsample.conv.weight" in vae_state_dict:

new_checkpoint[f"encoder.down_blocks.{i}.downsamplers.0.conv.weight"] = vae_state_dict.pop(

f"encoder.down.{i}.downsample.conv.weight"

)

new_checkpoint[f"encoder.down_blocks.{i}.downsamplers.0.conv.bias"] = vae_state_dict.pop(

f"encoder.down.{i}.downsample.conv.bias"

)

paths = renew_vae_resnet_paths(resnets)

meta_path = {"old": f"down.{i}.block", "new": f"down_blocks.{i}.resnets"}

assign_to_checkpoint(paths, new_checkpoint, vae_state_dict, additional_replacements=[meta_path], config=config)

mid_resnets = [key for key in vae_state_dict if "encoder.mid.block" in key]

num_mid_res_blocks = 2

for i in range(1, num_mid_res_blocks + 1):

resnets = [key for key in mid_resnets if f"encoder.mid.block_{i}" in key]

paths = renew_vae_resnet_paths(resnets)

meta_path = {"old": f"mid.block_{i}", "new": f"mid_block.resnets.{i - 1}"}

assign_to_checkpoint(paths, new_checkpoint, vae_state_dict, additional_replacements=[meta_path], config=config)

mid_attentions = [key for key in vae_state_dict if "encoder.mid.attn" in key]

paths = renew_vae_attention_paths(mid_attentions)

meta_path = {"old": "mid.attn_1", "new": "mid_block.attentions.0"}

assign_to_checkpoint(paths, new_checkpoint, vae_state_dict, additional_replacements=[meta_path], config=config)

conv_attn_to_linear(new_checkpoint)

for i in range(num_up_blocks):

block_id = num_up_blocks - 1 - i

resnets = [

key for key in up_blocks[block_id] if f"up.{block_id}" in key and f"up.{block_id}.upsample" not in key

]

if f"decoder.up.{block_id}.upsample.conv.weight" in vae_state_dict:

new_checkpoint[f"decoder.up_blocks.{i}.upsamplers.0.conv.weight"] = vae_state_dict[

f"decoder.up.{block_id}.upsample.conv.weight"

]

new_checkpoint[f"decoder.up_blocks.{i}.upsamplers.0.conv.bias"] = vae_state_dict[

f"decoder.up.{block_id}.upsample.conv.bias"

]

paths = renew_vae_resnet_paths(resnets)

meta_path = {"old": f"up.{block_id}.block", "new": f"up_blocks.{i}.resnets"}

assign_to_checkpoint(paths, new_checkpoint, vae_state_dict, additional_replacements=[meta_path], config=config)

mid_resnets = [key for key in vae_state_dict if "decoder.mid.block" in key]

num_mid_res_blocks = 2

for i in range(1, num_mid_res_blocks + 1):

resnets = [key for key in mid_resnets if f"decoder.mid.block_{i}" in key]

paths = renew_vae_resnet_paths(resnets)

meta_path = {"old": f"mid.block_{i}", "new": f"mid_block.resnets.{i - 1}"}

assign_to_checkpoint(paths, new_checkpoint, vae_state_dict, additional_replacements=[meta_path], config=config)

mid_attentions = [key for key in vae_state_dict if "decoder.mid.attn" in key]

paths = renew_vae_attention_paths(mid_attentions)

meta_path = {"old": "mid.attn_1", "new": "mid_block.attentions.0"}

assign_to_checkpoint(paths, new_checkpoint, vae_state_dict, additional_replacements=[meta_path], config=config)

conv_attn_to_linear(new_checkpoint)

return new_checkpoint

def vae_pt_to_vae_diffuser(

checkpoint_path: str,

output_path: str,

):

# Only support V1

r = requests.get(

" https://raw.githubusercontent.com/CompVis/stable-diffusion/main/configs/stable-diffusion/v1-inference.yaml"

)

io_obj = io.BytesIO(r.content)

original_config = OmegaConf.load(io_obj)

image_size = 512

device = "cuda" if torch.cuda.is_available() else "cpu"

checkpoint = torch.load(checkpoint_path, map_location=device)

# Convert the VAE model.

vae_config = create_vae_diffusers_config(original_config, image_size=image_size)

converted_vae_checkpoint = custom_convert_ldm_vae_checkpoint(checkpoint["state_dict"], vae_config)

vae = AutoencoderKL(**vae_config)

vae.load_state_dict(converted_vae_checkpoint)

vae.save_pretrained(output_path)

if __name__ == "__main__":

parser = argparse.ArgumentParser()

parser.add_argument("--vae_pt_path", default=None, type=str, required=True, help="Path to the VAE.pt to convert.")

parser.add_argument("--dump_path", default=None, type=str, required=True, help="Path to the VAE.pt to convert.")

args = parser.parse_args()

vae_pt_to_vae_diffuser(args.vae_pt_path, args.dump_path)