JackAILab commited on
Commit
35e9299
·
verified ·
1 Parent(s): c3e699e

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],