(self, rank, scale, ckpt_dir=None)
| 934 | self._init_lora(lora_rank, lora_scale, lora_weight_path) |
| 935 | |
| 936 | def _init_lora(self, rank, scale, ckpt_dir=None): |
| 937 | inject_trainable_lora_extended(self, target_replace_module=None, rank=rank, scale=scale) |
| 938 | |
| 939 | if ckpt_dir is not None: |
| 940 | with open(os.path.join(ckpt_dir, "latest")) as latest_file: |
| 941 | latest = latest_file.read().strip() |
| 942 | ckpt_path = os.path.join(ckpt_dir, latest, "mp_rank_00_model_states.pt") |
| 943 | print(f"loading lora from {ckpt_path}") |
| 944 | sd = th.load(ckpt_path)["module"] |
| 945 | sd = { |
| 946 | key[len("model.diffusion_model") :]: sd[key] for key in sd if key.startswith("model.diffusion_model") |
| 947 | } |
| 948 | self.load_state_dict(sd, strict=False) |
| 949 | |
| 950 | def _update_scale(self, scale): |
| 951 | update_lora_scale(self, scale) |
no test coverage detected