Update pipline_StableDiffusion_ConsistentID.py

pipline_StableDiffusion_ConsistentID.py

@@ -21,6 +21,7 @@ from attention import Consistent_IPAttProcessor, Consistent_AttProcessor, Facial
 ### We use the ckpt of 79999_iter.pth: https://drive.google.com/open?id=154JgKpzCPW82qINcVieuPH3fZ2e0P812
 ### Thanks for the open source of face-parsing model.
 from models.BiSeNet.model import BiSeNet
+bise_net_cp_path = hf_hub_download(repo_id="JackAILab/ConsistentID", filename="face_parsing.pth", repo_type="model")
 
 PipelineImageInput = Union[
     PIL.Image.Image,
@@ -31,7 +32,6 @@ PipelineImageInput = Union[
 
 ### Download the pretrained model from huggingface and put it locally, then place the model in a local directory and specify the directory location.
 class ConsistentIDStableDiffusionPipeline(StableDiffusionPipeline):
-    
     @validate_hf_hub_args
     def load_ConsistentID_model(
         self,
@@ -65,7 +65,7 @@ class ConsistentIDStableDiffusionPipeline(StableDiffusionPipeline):
         ### BiSeNet
         self.bise_net = BiSeNet(n_classes = 19)
         self.bise_net.cuda()
-        self.bise_net_cp='JackAILab/ConsistentID/face_parsing.pth' 
+        self.bise_net_cp=bise_net_cp_path
         self.bise_net.load_state_dict(torch.load(self.bise_net_cp))
         self.bise_net.eval()
         # Colors for all 20 parts
@@ -124,11 +124,12 @@ class ConsistentIDStableDiffusionPipeline(StableDiffusionPipeline):
             if weight_name.endswith(".safetensors"):
                 state_dict = {"id_encoder": {}, "lora_weights": {}}
                 with safe_open(model_file, framework="pt", device="cpu") as f:
+                    ### TODO safetensors add
                     for key in f.keys():
-                        if key.startswith("id_encoder."):
-                            state_dict["id_encoder"][key.replace("id_encoder.", "")] = f.get_tensor(key)
-                        elif key.startswith("lora_weights."):
-                            state_dict["lora_weights"][key.replace("lora_weights.", "")] = f.get_tensor(key)
+                        if key.startswith("FacialEncoder."):
+                            state_dict["FacialEncoder"][key.replace("FacialEncoder.", "")] = f.get_tensor(key)
+                        elif key.startswith("image_proj."):
+                            state_dict["image_proj"][key.replace("image_proj.", "")] = f.get_tensor(key)
             else:
                 state_dict = torch.load(model_file, map_location="cpu")
         else: