136 lines
5.0 KiB
Python
136 lines
5.0 KiB
Python
#encoding = utf8
|
|
import logging
|
|
import os
|
|
import queue
|
|
import time
|
|
from queue import Queue
|
|
from threading import Event, Thread
|
|
|
|
import numpy as np
|
|
import torch
|
|
|
|
from human_handler import AudioHandler
|
|
from utils import load_model, mirror_index, get_device, SyncQueue
|
|
|
|
logger = logging.getLogger(__name__)
|
|
current_file_path = os.path.dirname(os.path.abspath(__file__))
|
|
|
|
|
|
class AudioInferenceHandler(AudioHandler):
|
|
def __init__(self, context, handler):
|
|
super().__init__(context, handler)
|
|
|
|
self._mal_queue = Queue()
|
|
self._audio_queue = SyncQueue(context.batch_size * 2)
|
|
|
|
self._exit_event = Event()
|
|
self._run_thread = Thread(target=self.__on_run)
|
|
self._exit_event.set()
|
|
self._run_thread.start()
|
|
logger.info("AudioInferenceHandler init")
|
|
|
|
def on_handle(self, stream, type_):
|
|
if type_ == 1:
|
|
self._mal_queue.put(stream)
|
|
elif type_ == 0:
|
|
self._audio_queue.put(stream)
|
|
print('AudioInferenceHandler on_handle', type_, self._audio_queue.size())
|
|
|
|
def on_message(self, message):
|
|
super().on_message(message)
|
|
|
|
def __on_run(self):
|
|
wav2lip_path = os.path.join(current_file_path, '..', 'checkpoints', 'wav2lip.pth')
|
|
logger.info(f'AudioInferenceHandler init, path:{wav2lip_path}')
|
|
model = load_model(wav2lip_path)
|
|
logger.info("Model loaded")
|
|
|
|
face_list_cycle = self._context.face_list_cycle
|
|
|
|
length = len(face_list_cycle)
|
|
index = 0
|
|
count = 0
|
|
count_time = 0
|
|
logger.info('start inference')
|
|
|
|
device = get_device()
|
|
logger.info(f'use device:{device}')
|
|
|
|
while True:
|
|
if self._exit_event.is_set():
|
|
start_time = time.perf_counter()
|
|
batch_size = self._context.batch_size
|
|
try:
|
|
mel_batch = self._mal_queue.get(block=True, timeout=1)
|
|
except queue.Empty:
|
|
continue
|
|
|
|
print('origin mel_batch:', len(mel_batch))
|
|
is_all_silence = True
|
|
audio_frames = []
|
|
for _ in range(batch_size * 2):
|
|
frame, type_ = self._audio_queue.get()
|
|
print('AudioInferenceHandler type_', type_)
|
|
audio_frames.append((frame, type_))
|
|
if type_ == 0:
|
|
is_all_silence = False
|
|
if is_all_silence:
|
|
for i in range(batch_size):
|
|
self.on_next_handle((None, mirror_index(length, index), audio_frames[i * 2:i * 2 + 2]),
|
|
0)
|
|
index = index + 1
|
|
else:
|
|
logger.info('infer=======')
|
|
t = time.perf_counter()
|
|
img_batch = []
|
|
for i in range(batch_size):
|
|
idx = mirror_index(length, index + i)
|
|
face = face_list_cycle[idx]
|
|
img_batch.append(face)
|
|
|
|
print('orign img_batch:', len(img_batch), 'origin mel_batch:', len(mel_batch))
|
|
img_batch = np.asarray(img_batch)
|
|
mel_batch = np.asarray(mel_batch)
|
|
img_masked = img_batch.copy()
|
|
img_masked[:, face.shape[0] // 2:] = 0
|
|
|
|
img_batch = np.concatenate((img_masked, img_batch), axis=3) / 255.
|
|
mel_batch = np.reshape(mel_batch,
|
|
[len(mel_batch), mel_batch.shape[1], mel_batch.shape[2], 1])
|
|
|
|
img_batch = torch.FloatTensor(np.transpose(img_batch, (0, 3, 1, 2))).to(device)
|
|
mel_batch = torch.FloatTensor(np.transpose(mel_batch, (0, 3, 1, 2))).to(device)
|
|
print('img_batch:', img_batch.shape, 'mel_batch:', mel_batch.shape)
|
|
with torch.no_grad():
|
|
pred = model(mel_batch, img_batch)
|
|
|
|
pred = pred.cpu().numpy().transpose(0, 2, 3, 1) * 255.
|
|
|
|
count_time += (time.perf_counter() - t)
|
|
count += batch_size
|
|
|
|
if count >= 100:
|
|
logger.info(f"------actual avg infer fps:{count / count_time:.4f}")
|
|
count = 0
|
|
count_time = 0
|
|
|
|
for i, res_frame in enumerate(pred):
|
|
self.on_next_handle(
|
|
(res_frame, mirror_index(length, index), audio_frames[i * 2:i * 2 + 2]),
|
|
0)
|
|
index = index + 1
|
|
logger.info(f'total batch time: {time.perf_counter() - start_time}')
|
|
else:
|
|
time.sleep(1)
|
|
break
|
|
logger.info('AudioInferenceHandler inference processor stop')
|
|
|
|
def stop(self):
|
|
self._exit_event.clear()
|
|
self._run_thread.join()
|
|
|
|
def pause_talk(self):
|
|
pass
|
|
# self._audio_queue.clear()
|
|
# self._mal_queue.queue.clear()
|