| 1527 | self.init_from_ckpt(ckpt_path, ignore_keys) |
| 1528 | |
| 1529 | def init_from_ckpt(self, path, ignore_keys=list(), only_model=False): |
| 1530 | sd = torch.load(path, map_location="cpu") |
| 1531 | if "state_dict" in list(sd.keys()): |
| 1532 | sd = sd["state_dict"] |
| 1533 | keys = list(sd.keys()) |
| 1534 | for k in keys: |
| 1535 | for ik in ignore_keys: |
| 1536 | if k.startswith(ik): |
| 1537 | print("Deleting key {} from state_dict.".format(k)) |
| 1538 | del sd[k] |
| 1539 | |
| 1540 | # make it explicit, finetune by including extra input channels |
| 1541 | if exists(self.finetune_keys) and k in self.finetune_keys: |
| 1542 | new_entry = None |
| 1543 | for name, param in self.named_parameters(): |
| 1544 | if name in self.finetune_keys: |
| 1545 | print( |
| 1546 | f"modifying key '{name}' and keeping its original {self.keep_dims} (channels) dimensions only") |
| 1547 | new_entry = torch.zeros_like(param) # zero init |
| 1548 | assert exists(new_entry), 'did not find matching parameter to modify' |
| 1549 | new_entry[:, :self.keep_dims, ...] = sd[k] |
| 1550 | sd[k] = new_entry |
| 1551 | |
| 1552 | missing, unexpected = self.load_state_dict(sd, strict=False) if not only_model else self.model.load_state_dict( |
| 1553 | sd, strict=False) |
| 1554 | print(f"Restored from {path} with {len(missing)} missing and {len(unexpected)} unexpected keys") |
| 1555 | if len(missing) > 0: |
| 1556 | print(f"Missing Keys: {missing}") |
| 1557 | if len(unexpected) > 0: |
| 1558 | print(f"Unexpected Keys: {unexpected}") |
| 1559 | |
| 1560 | @torch.no_grad() |
| 1561 | def log_images(self, batch, N=8, n_row=4, sample=True, ddim_steps=200, ddim_eta=1., return_keys=None, |