yanze commited on
Commit
ca80b43
1 Parent(s): 3d99b24

Update pulid/pipeline_flux.py

Browse files
Files changed (1) hide show
  1. pulid/pipeline_flux.py +4 -1
pulid/pipeline_flux.py CHANGED
@@ -72,7 +72,7 @@ class PuLIDPipeline(nn.Module):
72
  name='antelopev2', root='.', providers=['CPUExecutionProvider']
73
  )
74
  self.app.prepare(ctx_id=0, det_size=(640, 640))
75
- self.handler_ante = insightface.model_zoo.get_model('models/antelopev2/glintr100.onnx')
76
  self.handler_ante.prepare(ctx_id=0)
77
 
78
  gc.collect()
@@ -117,6 +117,8 @@ class PuLIDPipeline(nn.Module):
117
  self.debug_img_list = []
118
  image_bgr = cv2.cvtColor(image, cv2.COLOR_RGB2BGR)
119
  # get antelopev2 embedding
 
 
120
  face_info = self.app.get(image_bgr)
121
  if len(face_info) > 0:
122
  face_info = sorted(face_info, key=lambda x: (x['bbox'][2] - x['bbox'][0]) * (x['bbox'][3] - x['bbox'][1]))[
@@ -142,6 +144,7 @@ class PuLIDPipeline(nn.Module):
142
  # incase insightface didn't detect face
143
  if id_ante_embedding is None:
144
  print('fail to detect face using insightface, extract embedding on align face')
 
145
  id_ante_embedding = self.handler_ante.get_feat(align_face)
146
 
147
  id_ante_embedding = torch.from_numpy(id_ante_embedding).to(self.device, self.weight_dtype)
 
72
  name='antelopev2', root='.', providers=['CPUExecutionProvider']
73
  )
74
  self.app.prepare(ctx_id=0, det_size=(640, 640))
75
+ self.handler_ante = insightface.model_zoo.get_model('models/antelopev2/glintr100.onnx', providers=['CPUExecutionProvider'])
76
  self.handler_ante.prepare(ctx_id=0)
77
 
78
  gc.collect()
 
117
  self.debug_img_list = []
118
  image_bgr = cv2.cvtColor(image, cv2.COLOR_RGB2BGR)
119
  # get antelopev2 embedding
120
+ for k in self.app.models.keys():
121
+ self.app.models[k].session.set_providers(['CUDAExecutionProvider'])
122
  face_info = self.app.get(image_bgr)
123
  if len(face_info) > 0:
124
  face_info = sorted(face_info, key=lambda x: (x['bbox'][2] - x['bbox'][0]) * (x['bbox'][3] - x['bbox'][1]))[
 
144
  # incase insightface didn't detect face
145
  if id_ante_embedding is None:
146
  print('fail to detect face using insightface, extract embedding on align face')
147
+ self.handler_ante.session.set_providers(['CUDAExecutionProvider'])
148
  id_ante_embedding = self.handler_ante.get_feat(align_face)
149
 
150
  id_ante_embedding = torch.from_numpy(id_ante_embedding).to(self.device, self.weight_dtype)