Sunbelt Computer Software

PL/B Language Development and Support

diffusers/src/diffusers/models/modeling_flax_pytorch_utils.py at main · anagri/diffusers · GitHub

Repository navigation

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

# 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.

""" PyTorch - Flax general utilities."""

import re

import jax.numpy as jnp

from flax.traverse_util import flatten_dict, unflatten_dict

from jax.random import PRNGKey

from ..utils import logging

logger = logging.get_logger(__name__)

def rename_key(key):

regex = r"\w+[.]\d+"

pats = re.findall(regex, key)

for pat in pats:

key = key.replace(pat, "_".join(pat.split(".")))

return key

#####################

# PyTorch => Flax #

#####################

# Adapted from https://github.com/huggingface/transformers/blob/c603c80f46881ae18b2ca50770ef65fa4033eacd/src/transformers/modeling_flax_pytorch_utils.py#L69

# and https://github.com/patil-suraj/stable-diffusion-jax/blob/main/stable_diffusion_jax/convert_diffusers_to_jax.py

def rename_key_and_reshape_tensor(pt_tuple_key, pt_tensor, random_flax_state_dict):

"""Rename PT weight names to corresponding Flax weight names and reshape tensor if necessary"""

# conv norm or layer norm

renamed_pt_tuple_key = pt_tuple_key[:-1] + ("scale",)

if (

any("norm" in str_ for str_ in pt_tuple_key)

and (pt_tuple_key[-1] == "bias")

and (pt_tuple_key[:-1] + ("bias",) not in random_flax_state_dict)

and (pt_tuple_key[:-1] + ("scale",) in random_flax_state_dict)

):

renamed_pt_tuple_key = pt_tuple_key[:-1] + ("scale",)

return renamed_pt_tuple_key, pt_tensor

elif pt_tuple_key[-1] in ["weight", "gamma"] and pt_tuple_key[:-1] + ("scale",) in random_flax_state_dict:

renamed_pt_tuple_key = pt_tuple_key[:-1] + ("scale",)

return renamed_pt_tuple_key, pt_tensor

# embedding

if pt_tuple_key[-1] == "weight" and pt_tuple_key[:-1] + ("embedding",) in random_flax_state_dict:

pt_tuple_key = pt_tuple_key[:-1] + ("embedding",)

return renamed_pt_tuple_key, pt_tensor

# conv layer

renamed_pt_tuple_key = pt_tuple_key[:-1] + ("kernel",)

if pt_tuple_key[-1] == "weight" and pt_tensor.ndim == 4:

pt_tensor = pt_tensor.transpose(2, 3, 1, 0)

return renamed_pt_tuple_key, pt_tensor

# linear layer

renamed_pt_tuple_key = pt_tuple_key[:-1] + ("kernel",)

if pt_tuple_key[-1] == "weight":

pt_tensor = pt_tensor.T

return renamed_pt_tuple_key, pt_tensor

# old PyTorch layer norm weight

renamed_pt_tuple_key = pt_tuple_key[:-1] + ("weight",)

if pt_tuple_key[-1] == "gamma":

return renamed_pt_tuple_key, pt_tensor

# old PyTorch layer norm bias

renamed_pt_tuple_key = pt_tuple_key[:-1] + ("bias",)

if pt_tuple_key[-1] == "beta":

return renamed_pt_tuple_key, pt_tensor

return pt_tuple_key, pt_tensor

def convert_pytorch_state_dict_to_flax(pt_state_dict, flax_model, init_key=42):

# Step 1: Convert pytorch tensor to numpy

pt_state_dict = {k: v.numpy() for k, v in pt_state_dict.items()}

# Step 2: Since the model is stateless, get random Flax params

random_flax_params = flax_model.init_weights(PRNGKey(init_key))

random_flax_state_dict = flatten_dict(random_flax_params)

flax_state_dict = {}

# Need to change some parameters name to match Flax names

for pt_key, pt_tensor in pt_state_dict.items():

renamed_pt_key = rename_key(pt_key)

pt_tuple_key = tuple(renamed_pt_key.split("."))

# Correctly rename weight parameters

flax_key, flax_tensor = rename_key_and_reshape_tensor(pt_tuple_key, pt_tensor, random_flax_state_dict)

if flax_key in random_flax_state_dict:

if flax_tensor.shape != random_flax_state_dict[flax_key].shape:

raise ValueError(

f"PyTorch checkpoint seems to be incorrect. Weight {pt_key} was expected to be of shape "

f"{random_flax_state_dict[flax_key].shape}, but is {flax_tensor.shape}."

)

# also add unexpected weight so that warning is thrown

flax_state_dict[flax_key] = jnp.asarray(flax_tensor)

return unflatten_dict(flax_state_dict)