Basis for different finetunas, such as inpainting or depth2image To disable finetuning mode, set finetune_keys to None
| 1498 | |
| 1499 | |
| 1500 | class LatentFinetuneDiffusion(LatentDiffusion): |
| 1501 | """ |
| 1502 | Basis for different finetunas, such as inpainting or depth2image |
| 1503 | To disable finetuning mode, set finetune_keys to None |
| 1504 | """ |
| 1505 | |
| 1506 | def __init__(self, |
| 1507 | concat_keys: tuple, |
| 1508 | finetune_keys=("model.diffusion_model.input_blocks.0.0.weight", |
| 1509 | "model_ema.diffusion_modelinput_blocks00weight" |
| 1510 | ), |
| 1511 | keep_finetune_dims=4, |
| 1512 | # if model was trained without concat mode before and we would like to keep these channels |
| 1513 | c_concat_log_start=None, # to log reconstruction of c_concat codes |
| 1514 | c_concat_log_end=None, |
| 1515 | *args, **kwargs |
| 1516 | ): |
| 1517 | ckpt_path = kwargs.pop("ckpt_path", None) |
| 1518 | ignore_keys = kwargs.pop("ignore_keys", list()) |
| 1519 | super().__init__(*args, **kwargs) |
| 1520 | self.finetune_keys = finetune_keys |
| 1521 | self.concat_keys = concat_keys |
| 1522 | self.keep_dims = keep_finetune_dims |
| 1523 | self.c_concat_log_start = c_concat_log_start |
| 1524 | self.c_concat_log_end = c_concat_log_end |
| 1525 | if exists(self.finetune_keys): assert exists(ckpt_path), 'can only finetune from a given checkpoint' |
| 1526 | if exists(ckpt_path): |
| 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: |
nothing calls this directly
no outgoing calls
no test coverage detected