Update app.py

app.py

@@ -14,7 +14,6 @@ from huggingface_hub import hf_hub_download
 ### 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
-from multiprocessing import Process, Queue, Manager
 
 # zero = torch.Tensor([0]).cuda()
 # print(zero.device) # <-- 'cpu' 🤔
@@ -37,27 +36,10 @@ pipe = ConsistentIDStableDiffusionPipeline.from_pretrained(
 ).to(device)
 
 ### Load other pretrained models
-@spaces.GPU
-def load_model(queue, bise_net_cp_path):
-    bise_net = BiSeNet(n_classes = 19)
-    bise_net.to(device)
-    bise_net.load_state_dict(torch.load(bise_net_cp_path))
-    bise_net.eval()
-    queue.put(bise_net)
 ## BiSenet
-bise_net_cp_path = hf_hub_download(repo_id="JackAILab/ConsistentID", filename="face_parsing.pth", repo_type="model")
-## Create a manager object
-with Manager() as manager:
-    # Create a queue to share data between processes
-    queue = manager.Queue()
-
-    # Create a new process and start it
-    p = Process(target=load_model, args=(queue, bise_net_cp_path))
-    p.start()
-
-    # Wait for the process to finish and get the result
-    p.join()
-    bise_net = queue.get()    
+bise_net_cp_path = hf_hub_download(repo_id="JackAILab/ConsistentID", filename="face_parsing.pth", local_dir="./checkpoints")
+bise_net = BiSeNet(n_classes = 19)
+bise_net.load_state_dict(torch.load(self.bise_net_cp, map_location="cpu"))
 
 ### Load consistentID_model checkpoint
 pipe.load_ConsistentID_model(
@@ -66,7 +48,7 @@ pipe.load_ConsistentID_model(
     subfolder="",
     weight_name=os.path.basename(consistentID_path),
     trigger_word="img",
-)     
+).to(device)     
 pipe.scheduler = EulerDiscreteScheduler.from_config(pipe.scheduler.config)