| 85 | return model |
| 86 | |
| 87 | def load_lora_weights(model, ckpt_path, strict=False, verbose=False): |
| 88 | ckpt = torch.load(ckpt_path, map_location="cpu") |
| 89 | if verbose: |
| 90 | print("-" * 40) |
| 91 | print(f"{'Loading weights':<16}: {ckpt_path}") |
| 92 | print(f"{'Model':<16}: {model.__class__.__name__}") |
| 93 | if "global_step" in ckpt: |
| 94 | print(f"{f'Global Step':<16}: {ckpt['global_step']:,}") |
| 95 | print(f"{'Strict':<16}: {'True' if strict else 'False'}") |
| 96 | print("-" * 40) |
| 97 | sd = ckpt["state_dict"] if 'state_dict' in ckpt else ckpt |
| 98 | # load only if this is a lora key |
| 99 | sd = {k: v for k, v in sd.items() if (".A." in k or ".B." in k)} |
| 100 | # exclude the first convolutional layer |
| 101 | sd = {k: v for k, v in sd.items() if not "input_blocks.0.0" in k} |
| 102 | missing, unexpected = model.load_state_dict(sd, strict=strict) |
| 103 | if len(missing) > 0: |
| 104 | warnings.warn(f"Load model weights - missing keys: {len(missing)}") |
| 105 | if verbose: |
| 106 | print(missing) |
| 107 | if len(unexpected) > 0: |
| 108 | warnings.warn(f"Load model weights - unexpected keys: {len(unexpected)}") |
| 109 | if verbose: |
| 110 | print(unexpected) |
| 111 | return model |
| 112 | |
| 113 | |
| 114 | def count_params(model): |