Update pipline_StableDiffusion_ConsistentID.py

pipline_StableDiffusion_ConsistentID.py

@@ -17,11 +17,6 @@ from functions import process_text_with_markers, masks_for_unique_values, fetch_
 from functions import ProjPlusModel, masks_for_unique_values
 from attention import Consistent_IPAttProcessor, Consistent_AttProcessor, FacialEncoder
 from huggingface_hub import hf_hub_download
-### Model can be imported from https://github.com/zllrunning/face-parsing.PyTorch?tab=readme-ov-file
-### 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,
@@ -36,6 +31,7 @@ class ConsistentIDStableDiffusionPipeline(StableDiffusionPipeline):
     def load_ConsistentID_model(
         self,
         pretrained_model_name_or_path_or_dict: Union[str, Dict[str, torch.Tensor]],
+        bise_net,
         weight_name: str,
         subfolder: str = '',
         trigger_word_ID: str = '<|image|>',
@@ -63,10 +59,11 @@ class ConsistentIDStableDiffusionPipeline(StableDiffusionPipeline):
         self.app.prepare(ctx_id=0, det_size=(640, 640))
 
         ### BiSeNet
-        self.bise_net = BiSeNet(n_classes = 19)
+        # self.bise_net = BiSeNet(n_classes = 19)
         # self.bise_net.cuda() # CUDA must not be initialized in the main process on Spaces with Stateless GPU environment
-        self.bise_net_cp=bise_net_cp_path
-        self.bise_net.load_state_dict(torch.load(self.bise_net_cp), map_location=torch.device('cpu'))
+        # self.bise_net_cp=bise_net_cp_path
+        # self.bise_net.load_state_dict(torch.load(self.bise_net_cp))
+        self.bise_net = bise_net # load from outside
         self.bise_net.eval()
         # Colors for all 20 parts
         self.part_colors = [[255, 0, 0], [255, 85, 0], [255, 170, 0],