diff --git a/args.py b/args.py index ea13ae3..368d1d1 100644 --- a/args.py +++ b/args.py @@ -24,7 +24,8 @@ parser.add_argument('--skip_model', action='store_true') parser.add_argument('--ifm', type=str) parser.add_argument('--anime4k', action='store_true') parser.add_argument('--alpha_split', action='store_true') -parser.add_argument('--cache', type=str, default='100mb') +parser.add_argument('--cache', type=str, default='256mb') +parser.add_argument('--gpu_cache', type=str, default='512mb') parser.add_argument('--simplify', type=int, default=1) args = parser.parse_args() args.output_w = int(args.output_size.split('x')[0]) @@ -33,4 +34,8 @@ if args.cache is not None: args.max_cache_len=int(convert_to_byte(args.cache)/262144) else: args.max_cache_len=0 +if args.gpu_cache is not None: + args.max_gpu_cache_len=int(convert_to_byte(args.gpu_cache)/589824) +else: + args.max_gpu_cache_len=0 if args.output_webcam is None and args.output_dir is None: args.debug = True \ No newline at end of file diff --git a/main.py b/main.py index 04bf565..36d7f49 100644 --- a/main.py +++ b/main.py @@ -165,6 +165,7 @@ class ModelClientProcess(Process): self.model_fps_number = Value('f', 0.0) self.gpu_fps_number = Value('f', 0.0) self.cache_hit_ratio = Value('f', 0.0) + self.gpu_cache_hit_ratio= Value('f', 0.0) def run(self): model = None @@ -302,22 +303,23 @@ class ModelClientProcess(Process): input_hash = hash(tuple(model_input)) cached = model_cache.get(input_hash) tot += 1 + mouth_eye_vector_c=[0.0]*27 if not cached is None: self.output_queue.put_nowait(cached) model_cache.move_to_end(input_hash) hit += 1 - self.cache_hit_ratio.value = hit / tot else: if args.perf: tic = time.perf_counter() for i in range(27): mouth_eye_vector[0, i] = model_input[i] + mouth_eye_vector_c[i]=model_input[i] for i in range(3): pose_vector[0, i] = model_input[i + 27] if model is None: output_image = input_image else: - output_image = model(input_image, mouth_eye_vector, pose_vector) + output_image = model(input_image, mouth_eye_vector, pose_vector, mouth_eye_vector_c,self.gpu_cache_hit_ratio) if args.perf: torch.cuda.synchronize() print("model", (time.perf_counter() - tic) * 1000) @@ -337,9 +339,10 @@ class ModelClientProcess(Process): if args.max_cache_len > 0: model_cache[input_hash] = postprocessed_image if len(model_cache) > args.max_cache_len: - model_cache.popitem(0) + model_cache.popitem(last=False) if args.debug: self.model_fps_number.value = model_fps() + self.cache_hit_ratio.value = hit / tot @torch.no_grad() @@ -639,7 +642,10 @@ def main(): cv2.putText(output_frame, str('IFM_FPS:%.1f' % client_process.ifm_fps_number.value), (0, 48), cv2.FONT_HERSHEY_PLAIN, 1, (0, 255, 0), 1) if args.max_cache_len > 0: - cv2.putText(output_frame, str('CACHED:%.1f%%' % (model_process.cache_hit_ratio.value * 100)), (0, 64), + cv2.putText(output_frame, str('MEMCACHED:%.1f%%' % (model_process.cache_hit_ratio.value * 100)), (0, 64), + cv2.FONT_HERSHEY_PLAIN, 1, (0, 255, 0), 1) + if args.max_gpu_cache_len > 0: + cv2.putText(output_frame, str('GPUCACHED:%.1f%%' % (model_process.gpu_cache_hit_ratio.value * 100)), (0, 80), cv2.FONT_HERSHEY_PLAIN, 1, (0, 255, 0), 1) cv2.imshow("frame", output_frame) # cv2.imshow("camera", debug_image) diff --git a/models.py b/models.py index e7bcf07..09f415b 100644 --- a/models.py +++ b/models.py @@ -6,18 +6,36 @@ from tha2.poser.modes.mode_20 import load_face_morpher, load_face_rotater, load_ from args import args +from collections import OrderedDict + class TalkingAnimeLight(nn.Module): def __init__(self): super(TalkingAnimeLight, self).__init__() self.face_morpher = load_face_morpher('pretrained/face_morpher.pt') self.two_algo_face_rotator = load_face_rotater('pretrained/two_algo_face_rotator.pt') self.combiner = load_combiner('pretrained/combiner.pt') + self.face_cache = OrderedDict() + self.tot = 0 + self.hit = 0 - def forward(self, image, mouth_eye_vector, pose_vector): + def forward(self, image, mouth_eye_vector, pose_vector, mouth_eye_vector_c, ratio=None): x = image.clone() if args.perf: tic=time.perf_counter() - mouth_eye_morp_image = self.face_morpher(image[:, :, 32:224, 32:224], mouth_eye_vector) + input_hash = hash(tuple(mouth_eye_vector_c)) + cached = self.face_cache.get(input_hash) + self.tot+=1 + if cached is None: + mouth_eye_morp_image = self.face_morpher(image[:, :, 32:224, 32:224], mouth_eye_vector) + self.face_cache[input_hash]=mouth_eye_morp_image.detach() + if len(self.face_cache) > args.max_gpu_cache_len: + self.face_cache.popitem(last=False) + else: + self.hit+=1 + mouth_eye_morp_image=cached + self.face_cache.move_to_end(input_hash) + if args.debug and ratio is not None: + ratio.value=self.hit / self.tot if args.perf: print(" - face_morpher",(time.perf_counter()-tic)*1000) tic=time.perf_counter()