Spaces:
Running
on
Zero
Running
on
Zero
Update pipline_StableDiffusion_ConsistentID.py
Browse files
pipline_StableDiffusion_ConsistentID.py
CHANGED
@@ -66,7 +66,7 @@ class ConsistentIDStableDiffusionPipeline(StableDiffusionPipeline):
|
|
66 |
self.bise_net = BiSeNet(n_classes = 19)
|
67 |
# self.bise_net.cuda() # CUDA must not be initialized in the main process on Spaces with Stateless GPU environment
|
68 |
self.bise_net_cp=bise_net_cp_path
|
69 |
-
self.bise_net.load_state_dict(torch.load(self.bise_net_cp))
|
70 |
self.bise_net.eval()
|
71 |
# Colors for all 20 parts
|
72 |
self.part_colors = [[255, 0, 0], [255, 85, 0], [255, 170, 0],
|
|
|
66 |
self.bise_net = BiSeNet(n_classes = 19)
|
67 |
# self.bise_net.cuda() # CUDA must not be initialized in the main process on Spaces with Stateless GPU environment
|
68 |
self.bise_net_cp=bise_net_cp_path
|
69 |
+
self.bise_net.load_state_dict(torch.load(self.bise_net_cp), map_location=torch.device('cpu'))
|
70 |
self.bise_net.eval()
|
71 |
# Colors for all 20 parts
|
72 |
self.part_colors = [[255, 0, 0], [255, 85, 0], [255, 170, 0],
|