(self, file_path="", state_dict={}, device="cuda", torch_dtype=torch.float16, **kwargs)
| 174 | |
| 175 | |
| 176 | def load(self, file_path="", state_dict={}, device="cuda", torch_dtype=torch.float16, **kwargs): |
| 177 | if len(state_dict) == 0: |
| 178 | state_dict = load_state_dict(file_path) |
| 179 | |
| 180 | # Load models with strict matching |
| 181 | keys_hash_with_shape = hash_state_dict_keys(state_dict, with_shape=True) |
| 182 | if keys_hash_with_shape in self.keys_hash_with_shape_dict: |
| 183 | model_names, model_classes, model_resource = self.keys_hash_with_shape_dict[keys_hash_with_shape] |
| 184 | loaded_model_names, loaded_models = load_model_from_single_file(state_dict, model_names, model_classes, model_resource, torch_dtype, device) |
| 185 | return loaded_model_names, loaded_models |
| 186 | |
| 187 | # Load models without strict matching |
| 188 | # (the shape of parameters may be inconsistent, and the state_dict_converter will modify the model architecture) |
| 189 | keys_hash = hash_state_dict_keys(state_dict, with_shape=False) |
| 190 | if keys_hash in self.keys_hash_dict: |
| 191 | model_names, model_classes, model_resource = self.keys_hash_dict[keys_hash] |
| 192 | loaded_model_names, loaded_models = load_model_from_single_file(state_dict, model_names, model_classes, model_resource, torch_dtype, device) |
| 193 | return loaded_model_names, loaded_models |
| 194 | |
| 195 | return loaded_model_names, loaded_models |
| 196 | |
| 197 | |
| 198 |
nothing calls this directly
no test coverage detected