(self, file_path, model_type=None, device=None, torch_dtype=None)
| 121 | self.model_name = [] |
| 122 | |
| 123 | def load_model(self, file_path, model_type=None, device=None, torch_dtype=None): |
| 124 | print(f"Loading models from: {file_path}") |
| 125 | if isinstance(file_path, list): |
| 126 | state_dict = {} |
| 127 | for path in file_path: |
| 128 | state_dict.update(load_state_dict(path)) |
| 129 | elif os.path.isfile(file_path): |
| 130 | state_dict = load_state_dict(file_path) |
| 131 | else: |
| 132 | state_dict = None |
| 133 | |
| 134 | if model_type=="vae": |
| 135 | model_names = ['video_vae'] |
| 136 | model_classes = [WanVideoVAE] |
| 137 | elif model_type=="text_encoder": |
| 138 | model_names = ['video_text_encoder'] |
| 139 | model_classes = [WanTextEncoder] |
| 140 | elif model_type=="dit": |
| 141 | model_names = ['video_dit', 'video_vace'] |
| 142 | model_classes = [WanModel, VaceWanModel] |
| 143 | |
| 144 | model_names, models = load_model_from_single_file(state_dict, model_names, model_classes, torch_dtype, device) |
| 145 | |
| 146 | for model_name, model in zip(model_names, models): |
| 147 | self.model.append(model) |
| 148 | self.model_path.append(file_path) |
| 149 | self.model_name.append(model_name) |
| 150 | |
| 151 | |
| 152 | def fetch_model(self, model_name): |
no test coverage detected