(state_dict, model_names, model_classes, model_resource, torch_dtype, device)
| 54 | |
| 55 | |
| 56 | def load_model_from_single_file(state_dict, model_names, model_classes, model_resource, torch_dtype, device): |
| 57 | loaded_model_names, loaded_models = [], [] |
| 58 | for model_name, model_class in zip(model_names, model_classes): |
| 59 | print(f" model_name: {model_name} model_class: {model_class.__name__}") |
| 60 | state_dict_converter = model_class.state_dict_converter() |
| 61 | if model_resource == "civitai": |
| 62 | state_dict_results = state_dict_converter.from_civitai(state_dict) |
| 63 | elif model_resource == "diffusers": |
| 64 | state_dict_results = state_dict_converter.from_diffusers(state_dict) |
| 65 | if isinstance(state_dict_results, tuple): |
| 66 | model_state_dict, extra_kwargs = state_dict_results |
| 67 | print(f" This model is initialized with extra kwargs: {extra_kwargs}") |
| 68 | else: |
| 69 | model_state_dict, extra_kwargs = state_dict_results, {} |
| 70 | torch_dtype = torch.float32 if extra_kwargs.get("upcast_to_float32", False) else torch_dtype |
| 71 | with init_weights_on_device(): |
| 72 | model = model_class(**extra_kwargs) |
| 73 | if hasattr(model, "eval"): |
| 74 | model = model.eval() |
| 75 | model.load_state_dict(model_state_dict, assign=True) |
| 76 | model = model.to(dtype=torch_dtype, device=device) |
| 77 | loaded_model_names.append(model_name) |
| 78 | loaded_models.append(model) |
| 79 | return loaded_model_names, loaded_models |
| 80 | |
| 81 | |
| 82 | def load_model_from_huggingface_folder(file_path, model_names, model_classes, torch_dtype, device): |
no test coverage detected