| 770 | self.load_from_ckpt(load_from_ckpt) |
| 771 | |
| 772 | def load_from_ckpt(self, ckpt_path): |
| 773 | input_ch = self.state_dict()["input_blocks.0.0.weight"].shape[1] |
| 774 | assert input_ch >= 4 and input_ch // 4 * 4 == input_ch, "Input channels must be at a multiplier 4 to load from SD ckpt" |
| 775 | output_ch = self.state_dict()["out.2.weight"].shape[0] |
| 776 | assert output_ch >= 4 and output_ch // 4 * 4 == output_ch, "Output channels must be at a multiplier 4 to load from SD ckpt" |
| 777 | sd = th.load(ckpt_path) |
| 778 | sd_ = {} |
| 779 | for k,v in sd["state_dict"].items(): |
| 780 | if k.startswith("model.diffusion_model"): |
| 781 | sd_[k.replace("model.diffusion_model.", "")] = v |
| 782 | |
| 783 | if input_ch > 4: |
| 784 | # Scaling for input channels so that the gradients are not too large |
| 785 | scale = input_ch // 4 |
| 786 | sd_["input_blocks.0.0.weight"] = sd_["input_blocks.0.0.weight"] / scale |
| 787 | sd_["input_blocks.0.0.weight"] = sd_["input_blocks.0.0.weight"].repeat(1, scale, 1, 1) |
| 788 | |
| 789 | if output_ch > 4: |
| 790 | # No scaling for output channels |
| 791 | scale = output_ch // 4 |
| 792 | sd_["out.2.weight"] = sd_["out.2.weight"].repeat(scale, 1, 1, 1) |
| 793 | sd_["out.2.bias"] = sd_["out.2.bias"].repeat(scale) |
| 794 | |
| 795 | missing, unexpected = self.load_state_dict(sd_, strict=False) |
| 796 | |
| 797 | if len(missing) > 0: |
| 798 | print(f"Load model weights - missing keys: {len(missing)}") |
| 799 | print(missing) |
| 800 | if len(unexpected) > 0: |
| 801 | print(f"Load model weights - unexpected keys: {len(unexpected)}") |
| 802 | print(unexpected) |
| 803 | |
| 804 | |
| 805 | def convert_to_fp16(self): |