This commit is contained in:
lllyasviel
2024-01-18 04:40:29 -08:00
parent 75f5f0da01
commit 3ba5754cc7
15 changed files with 1265 additions and 57 deletions

View File

@@ -65,6 +65,12 @@ class BASE:
replace_prefix = {"": "cond_stage_model."}
return utils.state_dict_prefix_replace(state_dict, replace_prefix)
def process_clip_vision_state_dict_for_saving(self, state_dict):
replace_prefix = {}
if self.clip_vision_prefix is not None:
replace_prefix[""] = self.clip_vision_prefix
return utils.state_dict_prefix_replace(state_dict, replace_prefix)
def process_unet_state_dict_for_saving(self, state_dict):
replace_prefix = {"": "model.diffusion_model."}
return utils.state_dict_prefix_replace(state_dict, replace_prefix)