Spaces:
Running
on
Zero
Running
on
Zero
Update pulid/pipeline_flux.py
Browse files- 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)
|