| 141 | |
| 142 | |
| 143 | def load_torch_file(ckpt, safe_load=False, device=None, return_metadata=False): |
| 144 | if device is None: |
| 145 | device = torch.device("cpu") |
| 146 | metadata = None |
| 147 | if ckpt.lower().endswith(".safetensors") or ckpt.lower().endswith(".sft"): |
| 148 | try: |
| 149 | with safetensors.safe_open(ckpt, framework="pt", device=device.type) as f: |
| 150 | sd = {} |
| 151 | for k in f.keys(): |
| 152 | sd[k] = f.get_tensor(k) |
| 153 | if return_metadata: |
| 154 | metadata = f.metadata() |
| 155 | except Exception as e: |
| 156 | if len(e.args) > 0: |
| 157 | message = e.args[0] |
| 158 | if "HeaderTooLarge" in message: |
| 159 | raise ValueError("{}\n\nFile path: {}\n\nThe safetensors file is corrupt or invalid. Make sure this is actually a safetensors file and not a ckpt or pt or other filetype.".format(message, ckpt)) |
| 160 | if "MetadataIncompleteBuffer" in message: |
| 161 | raise ValueError("{}\n\nFile path: {}\n\nThe safetensors file is corrupt/incomplete. Check the file size and make sure you have copied/downloaded it correctly.".format(message, ckpt)) |
| 162 | raise e |
| 163 | else: |
| 164 | if safe_load or ALWAYS_SAFE_LOAD: |
| 165 | pl_sd = torch.load(ckpt, map_location=device, weights_only=True) |
| 166 | else: |
| 167 | pl_sd = torch.load(ckpt, map_location=device, pickle_module=comfy.checkpoint_pickle) |
| 168 | if "global_step" in pl_sd: |
| 169 | logging.debug(f"Global Step: {pl_sd['global_step']}") |
| 170 | if "state_dict" in pl_sd: |
| 171 | sd = pl_sd["state_dict"] |
| 172 | else: |
| 173 | if len(pl_sd) == 1: |
| 174 | key = list(pl_sd.keys())[0] |
| 175 | sd = pl_sd[key] |
| 176 | if not isinstance(sd, dict): |
| 177 | sd = pl_sd |
| 178 | else: |
| 179 | sd = pl_sd |
| 180 | return (sd, metadata) if return_metadata else sd |
| 181 | |
| 182 | def standardize_lora_key_format(lora_sd): |
| 183 | new_sd = {} |