Update app.py

app.py

@@ -10,6 +10,10 @@ from diffusers.utils import load_image
 from diffusers import EulerDiscreteScheduler
 from pipline_StableDiffusion_ConsistentID import ConsistentIDStableDiffusionPipeline
 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
 
 zero = torch.Tensor([0]).cuda()
 print(zero.device) # <-- 'cpu' 🤔
@@ -30,9 +34,17 @@ pipe = ConsistentIDStableDiffusionPipeline.from_pretrained(
     variant="fp16"
 ).to(device)
 
+### Load other pretrained models
+## BiSenet
+bise_net = BiSeNet(n_classes = 19)
+bise_net.cuda() # CUDA must not be initialized in the main process on Spaces with Stateless GPU environment
+bise_net_cp_path = hf_hub_download(repo_id="JackAILab/ConsistentID", filename="face_parsing.pth", repo_type="model")
+bise_net.load_state_dict(torch.load(bise_net_cp_path))
+
 ### Load consistentID_model checkpoint
 pipe.load_ConsistentID_model(
     os.path.dirname(consistentID_path),
+    bise_net,
     subfolder="",
     weight_name=os.path.basename(consistentID_path),
     trigger_word="img",