(self, model_config, model_type)
| 142 | return filtered_checkpoint |
| 143 | |
| 144 | def load_model(self, model_config, model_type): |
| 145 | |
| 146 | device = get_device() |
| 147 | |
| 148 | if model_type == 'stitching_retargeting_module': |
| 149 | ckpt_path = os.path.join(get_model_dir("liveportrait"), "retargeting_models", model_type + ".pth") |
| 150 | else: |
| 151 | ckpt_path = os.path.join(get_model_dir("liveportrait"), "base_models", model_type + ".pth") |
| 152 | |
| 153 | is_safetensors = None |
| 154 | if os.path.isfile(ckpt_path) == False: |
| 155 | is_safetensors = True |
| 156 | ckpt_path = os.path.join(get_model_dir("liveportrait"), model_type + ".safetensors") |
| 157 | if os.path.isfile(ckpt_path) == False: |
| 158 | self.download_model(ckpt_path, |
| 159 | "https://huggingface.co/Kijai/LivePortrait_safetensors/resolve/main/" + model_type + ".safetensors") |
| 160 | model_params = model_config['model_params'][f'{model_type}_params'] |
| 161 | if model_type == 'appearance_feature_extractor': |
| 162 | model = AppearanceFeatureExtractor(**model_params).to(device) |
| 163 | elif model_type == 'motion_extractor': |
| 164 | model = MotionExtractor(**model_params).to(device) |
| 165 | elif model_type == 'warping_module': |
| 166 | model = WarpingNetwork(**model_params).to(device) |
| 167 | elif model_type == 'spade_generator': |
| 168 | model = SPADEDecoder(**model_params).to(device) |
| 169 | elif model_type == 'stitching_retargeting_module': |
| 170 | # Special handling for stitching and retargeting module |
| 171 | config = model_config['model_params']['stitching_retargeting_module_params'] |
| 172 | checkpoint = comfy.utils.load_torch_file(ckpt_path) |
| 173 | |
| 174 | stitcher = StitchingRetargetingNetwork(**config.get('stitching')) |
| 175 | if is_safetensors: |
| 176 | stitcher.load_state_dict(self.filter_for_model(checkpoint, 'retarget_shoulder')) |
| 177 | else: |
| 178 | stitcher.load_state_dict(self.remove_ddp_dumplicate_key(checkpoint['retarget_shoulder'])) |
| 179 | stitcher = stitcher.to(device) |
| 180 | stitcher.eval() |
| 181 | |
| 182 | return { |
| 183 | 'stitching': stitcher, |
| 184 | } |
| 185 | else: |
| 186 | raise ValueError(f"Unknown model type: {model_type}") |
| 187 | |
| 188 | |
| 189 | model.load_state_dict(comfy.utils.load_torch_file(ckpt_path)) |
| 190 | model.eval() |
| 191 | return model |
| 192 | |
| 193 | def load_models(self): |
| 194 | model_path = get_model_dir("liveportrait") |
no test coverage detected