Update app.py

app.py

@@ -14,6 +14,7 @@ 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
 
 zero = torch.Tensor([0]).cuda()
 print(zero.device) # <-- 'cpu' 🤔
@@ -36,10 +37,15 @@ pipe = ConsistentIDStableDiffusionPipeline.from_pretrained(
 
 ### 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))
+# Create a queue to share data between processes
+queue = 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()
 
 ### Load consistentID_model checkpoint
 pipe.load_ConsistentID_model(
@@ -51,6 +57,16 @@ pipe.load_ConsistentID_model(
 )     
 pipe.scheduler = EulerDiscreteScheduler.from_config(pipe.scheduler.config)
 
+
+@spaces.GPU
+def load_model(queue, bise_net_cp_path):
+    device = torch.device('cuda')
+    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)
+
 @spaces.GPU
 def process(inputImage,prompt,negative_prompt):