| 109 | |
| 110 | |
| 111 | class ModelManager: |
| 112 | def __init__( |
| 113 | self, |
| 114 | torch_dtype = torch.float16, |
| 115 | device = "cuda", |
| 116 | ): |
| 117 | self.torch_dtype = torch_dtype |
| 118 | self.device = device |
| 119 | self.model = [] |
| 120 | self.model_path = [] |
| 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): |
| 153 | fetched_models = [] |
| 154 | for model, model_path, model_name_ in zip(self.model, self.model_path, self.model_name): |
| 155 | if model_name == model_name_: |
| 156 | fetched_models.append(model) |
| 157 | if len(fetched_models) == 0: |
| 158 | print(f"No {model_name} models available.") |
| 159 | return None |
| 160 | if len(fetched_models) > 1: |
| 161 | raise ValueError( |
| 162 | f"Expected one {model_name} model, but found {len(fetched_models)}." |
| 163 | ) |
| 164 | return fetched_models[0] |