| 181 | return args |
| 182 | |
| 183 | def strip_prefix(state_dict): |
| 184 | # prefix for mixed state dict |
| 185 | prefix = None |
| 186 | for pfx in ["model.diffusion_model.", "model."]: |
| 187 | if any([x.startswith(pfx) for x in state_dict.keys()]): |
| 188 | prefix = pfx |
| 189 | break |
| 190 | |
| 191 | # prefix for uniform state dict |
| 192 | if prefix is None: |
| 193 | for pfx in ["net."]: |
| 194 | if all([x.startswith(pfx) for x in state_dict.keys()]): |
| 195 | prefix = pfx |
| 196 | break |
| 197 | |
| 198 | # strip prefix if found |
| 199 | if prefix is not None: |
| 200 | logging.info(f"State dict prefix found: '{prefix}'") |
| 201 | sd = {} |
| 202 | for k, v in state_dict.items(): |
| 203 | if prefix not in k: |
| 204 | continue |
| 205 | k = k.replace(prefix, "") |
| 206 | sd[k] = v |
| 207 | else: |
| 208 | logging.debug("State dict has no prefix") |
| 209 | sd = state_dict |
| 210 | |
| 211 | return sd |
| 212 | |
| 213 | def load_state_dict(path): |
| 214 | if any(path.endswith(x) for x in [".ckpt", ".pt", ".bin", ".pth"]): |