(old_transformer, model_256)
| 116 | |
| 117 | |
| 118 | def make_transformer(old_transformer, model_256): |
| 119 | args = dict(old_transformer.config) |
| 120 | force_down_up_sample = args["force_down_up_sample"] |
| 121 | |
| 122 | signature = inspect.signature(UVit2DModel.__init__) |
| 123 | |
| 124 | args_ = { |
| 125 | "downsample": force_down_up_sample, |
| 126 | "upsample": force_down_up_sample, |
| 127 | "block_out_channels": args["block_out_channels"][0], |
| 128 | "sample_size": 16 if model_256 else 32, |
| 129 | } |
| 130 | |
| 131 | for s in list(signature.parameters.keys()): |
| 132 | if s in ["self", "downsample", "upsample", "sample_size", "block_out_channels"]: |
| 133 | continue |
| 134 | |
| 135 | args_[s] = args[s] |
| 136 | |
| 137 | new_transformer = UVit2DModel(**args_) |
| 138 | new_transformer.to(device) |
| 139 | |
| 140 | new_transformer.set_attn_processor(AttnProcessor()) |
| 141 | |
| 142 | state_dict = old_transformer.state_dict() |
| 143 | |
| 144 | state_dict["cond_embed.linear_1.weight"] = state_dict.pop("cond_embed.0.weight") |
| 145 | state_dict["cond_embed.linear_2.weight"] = state_dict.pop("cond_embed.2.weight") |
| 146 | |
| 147 | for i in range(22): |
| 148 | state_dict[f"transformer_layers.{i}.norm1.norm.weight"] = state_dict.pop( |
| 149 | f"transformer_layers.{i}.attn_layer_norm.weight" |
| 150 | ) |
| 151 | state_dict[f"transformer_layers.{i}.norm1.linear.weight"] = state_dict.pop( |
| 152 | f"transformer_layers.{i}.self_attn_adaLN_modulation.mapper.weight" |
| 153 | ) |
| 154 | |
| 155 | state_dict[f"transformer_layers.{i}.attn1.to_q.weight"] = state_dict.pop( |
| 156 | f"transformer_layers.{i}.attention.query.weight" |
| 157 | ) |
| 158 | state_dict[f"transformer_layers.{i}.attn1.to_k.weight"] = state_dict.pop( |
| 159 | f"transformer_layers.{i}.attention.key.weight" |
| 160 | ) |
| 161 | state_dict[f"transformer_layers.{i}.attn1.to_v.weight"] = state_dict.pop( |
| 162 | f"transformer_layers.{i}.attention.value.weight" |
| 163 | ) |
| 164 | state_dict[f"transformer_layers.{i}.attn1.to_out.0.weight"] = state_dict.pop( |
| 165 | f"transformer_layers.{i}.attention.out.weight" |
| 166 | ) |
| 167 | |
| 168 | state_dict[f"transformer_layers.{i}.norm2.norm.weight"] = state_dict.pop( |
| 169 | f"transformer_layers.{i}.crossattn_layer_norm.weight" |
| 170 | ) |
| 171 | state_dict[f"transformer_layers.{i}.norm2.linear.weight"] = state_dict.pop( |
| 172 | f"transformer_layers.{i}.cross_attn_adaLN_modulation.mapper.weight" |
| 173 | ) |
| 174 | |
| 175 | state_dict[f"transformer_layers.{i}.attn2.to_q.weight"] = state_dict.pop( |
no test coverage detected
searching dependent graphs…