1- import os
2- import sys
3-
4- import numpy as np
5- import torch
6- from diffusers .utils import load_image
7- # pip install insightface==0.7.3
8- from insightface .app import FaceAnalysis
9- from insightface .data import get_image as ins_get_image
10- from safetensors .torch import save_file
11-
12- ###
13- # https://github.com/cubiq/ComfyUI_IPAdapter_plus/issues/165#issue-2055829543
14- ###
15- class FaceAnalysis2 (FaceAnalysis ):
16- # NOTE: allows setting det_size for each detection call.
17- # the model allows it but the wrapping code from insightface
18- # doesn't show it, and people end up loading duplicate models
19- # for different sizes where there is absolutely no need to
20- def get (self , img , max_num = 0 , det_size = (640 , 640 )):
21- if det_size is not None :
22- self .det_model .input_size = det_size
23-
24- return super ().get (img , max_num )
25-
26- def analyze_faces (face_analysis : FaceAnalysis , img_data : np .ndarray , det_size = (640 , 640 )):
27- # NOTE: try detect faces, if no faces detected, lower det_size until it does
28- detection_sizes = [None ] + [(size , size ) for size in range (640 , 256 , - 64 )] + [(256 , 256 )]
29-
30- for size in detection_sizes :
31- faces = face_analysis .get (img_data , det_size = size )
32- if len (faces ) > 0 :
33- return faces
34-
35- return []
36-
37- if __name__ == "__main__" :
38- #face_detector = FaceAnalysis2(providers=['CUDAExecutionProvider'], allowed_modules=['detection', 'recognition'])
39- face_detector = FaceAnalysis2 (providers = ['CPUExecutionProvider' ], allowed_modules = ['detection' , 'recognition' ])
40- face_detector .prepare (ctx_id = 0 , det_size = (640 , 640 ))
41- #input_folder_name = './scarletthead_woman'
42- input_folder_name = sys .argv [1 ]
43- image_basename_list = os .listdir (input_folder_name )
44- image_path_list = sorted ([os .path .join (input_folder_name , basename ) for basename in image_basename_list ])
45-
46- input_id_images = []
47- for image_path in image_path_list :
48- input_id_images .append (load_image (image_path ))
49-
50- id_embed_list = []
51-
52- for img in input_id_images :
53- img = np .array (img )
54- img = img [:, :, ::- 1 ]
55- faces = analyze_faces (face_detector , img )
56- if len (faces ) > 0 :
57- id_embed_list .append (torch .from_numpy ((faces [0 ]['embedding' ])))
58-
59- if len (id_embed_list ) == 0 :
60- raise ValueError (f"No face detected in input image pool" )
61-
62- id_embeds = torch .stack (id_embed_list )
63-
64- # for r in id_embeds:
65- # print(r)
66- # #torch.save(id_embeds, input_folder_name+'/id_embeds.pt');
67- # weights = dict()
68- # weights["id_embeds"] = id_embeds
69- # save_file(weights, input_folder_name+'/id_embeds.safetensors')
70-
71- binary_data = id_embeds .numpy ().tobytes ()
72- two = 4
73- zero = 0
74- one = 1
75- tensor_name = "id_embeds"
76- # Write binary data to a file
77- with open (input_folder_name + '/id_embeds.bin' , "wb" ) as f :
78- f .write (two .to_bytes (4 , byteorder = 'little' ))
79- f .write ((len (tensor_name )).to_bytes (4 , byteorder = 'little' ))
80- f .write (zero .to_bytes (4 , byteorder = 'little' ))
81- f .write ((id_embeds .shape [1 ]).to_bytes (4 , byteorder = 'little' ))
82- f .write ((id_embeds .shape [0 ]).to_bytes (4 , byteorder = 'little' ))
83- f .write (one .to_bytes (4 , byteorder = 'little' ))
84- f .write (one .to_bytes (4 , byteorder = 'little' ))
85- f .write (tensor_name .encode ('ascii' ))
86- f .write (binary_data )
87-
1+ import os
2+ import sys
3+
4+ import numpy as np
5+ import torch
6+ from diffusers .utils import load_image
7+ # pip install insightface==0.7.3
8+ from insightface .app import FaceAnalysis
9+ from insightface .data import get_image as ins_get_image
10+ from safetensors .torch import save_file
11+
12+ ###
13+ # https://github.com/cubiq/ComfyUI_IPAdapter_plus/issues/165#issue-2055829543
14+ ###
15+ class FaceAnalysis2 (FaceAnalysis ):
16+ # NOTE: allows setting det_size for each detection call.
17+ # the model allows it but the wrapping code from insightface
18+ # doesn't show it, and people end up loading duplicate models
19+ # for different sizes where there is absolutely no need to
20+ def get (self , img , max_num = 0 , det_size = (640 , 640 )):
21+ if det_size is not None :
22+ self .det_model .input_size = det_size
23+
24+ return super ().get (img , max_num )
25+
26+ def analyze_faces (face_analysis : FaceAnalysis , img_data : np .ndarray , det_size = (640 , 640 )):
27+ # NOTE: try detect faces, if no faces detected, lower det_size until it does
28+ detection_sizes = [None ] + [(size , size ) for size in range (640 , 256 , - 64 )] + [(256 , 256 )]
29+
30+ for size in detection_sizes :
31+ faces = face_analysis .get (img_data , det_size = size )
32+ if len (faces ) > 0 :
33+ return faces
34+
35+ return []
36+
37+ if __name__ == "__main__" :
38+ #face_detector = FaceAnalysis2(providers=['CUDAExecutionProvider'], allowed_modules=['detection', 'recognition'])
39+ face_detector = FaceAnalysis2 (providers = ['CPUExecutionProvider' ], allowed_modules = ['detection' , 'recognition' ])
40+ face_detector .prepare (ctx_id = 0 , det_size = (640 , 640 ))
41+ #input_folder_name = './scarletthead_woman'
42+ input_folder_name = sys .argv [1 ]
43+ image_basename_list = os .listdir (input_folder_name )
44+ image_path_list = sorted ([os .path .join (input_folder_name , basename ) for basename in image_basename_list ])
45+
46+ input_id_images = []
47+ for image_path in image_path_list :
48+ input_id_images .append (load_image (image_path ))
49+
50+ id_embed_list = []
51+
52+ for img in input_id_images :
53+ img = np .array (img )
54+ img = img [:, :, ::- 1 ]
55+ faces = analyze_faces (face_detector , img )
56+ if len (faces ) > 0 :
57+ id_embed_list .append (torch .from_numpy ((faces [0 ]['embedding' ])))
58+
59+ if len (id_embed_list ) == 0 :
60+ raise ValueError (f"No face detected in input image pool" )
61+
62+ id_embeds = torch .stack (id_embed_list )
63+
64+ # for r in id_embeds:
65+ # print(r)
66+ # #torch.save(id_embeds, input_folder_name+'/id_embeds.pt');
67+ # weights = dict()
68+ # weights["id_embeds"] = id_embeds
69+ # save_file(weights, input_folder_name+'/id_embeds.safetensors')
70+
71+ binary_data = id_embeds .numpy ().tobytes ()
72+ two = 4
73+ zero = 0
74+ one = 1
75+ tensor_name = "id_embeds"
76+ # Write binary data to a file
77+ with open (input_folder_name + '/id_embeds.bin' , "wb" ) as f :
78+ f .write (two .to_bytes (4 , byteorder = 'little' ))
79+ f .write ((len (tensor_name )).to_bytes (4 , byteorder = 'little' ))
80+ f .write (zero .to_bytes (4 , byteorder = 'little' ))
81+ f .write ((id_embeds .shape [1 ]).to_bytes (4 , byteorder = 'little' ))
82+ f .write ((id_embeds .shape [0 ]).to_bytes (4 , byteorder = 'little' ))
83+ f .write (one .to_bytes (4 , byteorder = 'little' ))
84+ f .write (one .to_bytes (4 , byteorder = 'little' ))
85+ f .write (tensor_name .encode ('ascii' ))
86+ f .write (binary_data )
87+
8888
0 commit comments