399 lines
14 KiB
Python
399 lines
14 KiB
Python
import argparse
|
||
import os
|
||
import warnings
|
||
import torch
|
||
from diffusers import AutoencoderKL, DDIMScheduler
|
||
from omegaconf import OmegaConf
|
||
from torch import nn
|
||
import cv2
|
||
import numpy as np
|
||
from typing import List
|
||
from einops import rearrange
|
||
import librosa
|
||
import soundfile as sf
|
||
import torchaudio
|
||
import imageio
|
||
import subprocess
|
||
from denoiser import pretrained
|
||
from denoiser.dsp import convert_audio
|
||
|
||
from .src.pipelines.face_animate import FaceAnimatePipeline
|
||
from .src.models.unet_2d_condition import UNet2DConditionModel
|
||
from .src.models.unet_3d_echo import EchoUNet3DConditionModel
|
||
from .src.models.whisper.audio2feature import load_audio_model
|
||
from .image_processor import ImageProcessor
|
||
from .src.utils.util import save_video, filter_non_none
|
||
|
||
|
||
warnings.filterwarnings("ignore")
|
||
|
||
class Net(nn.Module):
|
||
"""
|
||
The Net class combines all the necessary modules for the inference process.
|
||
|
||
Args:
|
||
reference_unet (UNet2DConditionModel): The UNet2DConditionModel used as a reference for inference.
|
||
denoising_unet (UNet3DConditionModel): The UNet3DConditionModel used for denoising the input audio.
|
||
"""
|
||
def __init__(
|
||
self,
|
||
reference_unet: UNet2DConditionModel,
|
||
denoising_unet: EchoUNet3DConditionModel,
|
||
):
|
||
super().__init__()
|
||
self.reference_unet = reference_unet
|
||
self.denoising_unet = denoising_unet
|
||
|
||
def forward(self,):
|
||
"""
|
||
empty function to override abstract function of nn Module
|
||
"""
|
||
|
||
def get_modules(self):
|
||
"""
|
||
Simple method to avoid too-few-public-methods pylint error
|
||
"""
|
||
return {
|
||
"reference_unet": self.reference_unet,
|
||
"denoising_unet": self.denoising_unet,
|
||
}
|
||
|
||
|
||
def create_mask(mask_path, mask_512_path, clip_length):
|
||
mask = 1 - (cv2.imread(mask_path) / 255.)
|
||
(h, w, c) = mask.shape
|
||
new_mask = np.zeros((h, w, c+1))
|
||
new_mask[:, :, :c] = mask
|
||
new_mask[:, :, c] = mask[:, :, c-1]
|
||
|
||
mask = torch.FloatTensor(new_mask).permute(2,0,1).unsqueeze(1)
|
||
|
||
return mask
|
||
|
||
|
||
def create_mask_alpha(w, h, coordinate):
|
||
(left, top, right, bottom) = coordinate
|
||
|
||
mask = np.ones((h,w,3))
|
||
mask[top:bottom, left:right] = 0
|
||
|
||
pad = min(int((right-left)/15), int((bottom-top)/15)) # 100/15=6.6% for padding
|
||
|
||
for idx in range(left, right):
|
||
mask[bottom-pad:bottom, idx] = np.array([np.linspace(0,1,pad)]*3).transpose(1,0)
|
||
mask[top:top+pad, idx] = np.array([np.linspace(1,0,pad)]*3).transpose(1,0)
|
||
|
||
for idx in range(top, bottom):
|
||
mask[idx, right-pad:right] = np.array([np.linspace(0,1,pad)]*3).transpose(1,0)
|
||
mask[idx, left:left+pad] = np.array([np.linspace(1,0,pad)]*3).transpose(1,0)
|
||
|
||
for idx, start, end in zip(range(top, top+pad), np.linspace(1,0.7, pad), np.linspace(0.7,0, pad)):
|
||
mask[idx, left:left+pad] = np.array([np.linspace(start,end,pad)]*3).transpose(1,0)
|
||
mask[idx, right-pad:right] = np.array([np.linspace(end,start,pad)]*3).transpose(1,0)
|
||
|
||
for idx, start, end in zip(range(bottom-pad, bottom), np.linspace(0.7,1, pad), np.linspace(0,0.7, pad)):
|
||
mask[idx, left:left+pad] = np.array([np.linspace(start,end,pad)]*3).transpose(1,0)
|
||
mask[idx, right-pad:right] = np.array([np.linspace(end,start,pad)]*3).transpose(1,0)
|
||
|
||
return mask
|
||
|
||
def process_audio(input_audio_path, output_audio_path):
|
||
wav,_ = librosa.load(input_audio_path, sr=16000)
|
||
tmp_path = output_audio_path[:-4]+"_16000.wav"
|
||
sf.write(tmp_path, wav, 16000)
|
||
|
||
model = pretrained.dns64().cuda()
|
||
wav, sr = torchaudio.load(input_audio_path)
|
||
wav = convert_audio(wav.cuda(), sr, model.sample_rate, model.chin)
|
||
|
||
with torch.no_grad():
|
||
denoised = model(wav[None])[0]
|
||
|
||
sf.write(output_audio_path, denoised.data.cpu().numpy()[0], 16000)
|
||
|
||
return
|
||
|
||
img_size = (512,512)
|
||
|
||
|
||
def mem_gpu():
|
||
# Total memory
|
||
total = torch.cuda.get_device_properties(0).total_memory
|
||
# Allocated memory
|
||
allocated = torch.cuda.memory_allocated(0)
|
||
# Cached memory (used for optimization)
|
||
cached = torch.cuda.memory_reserved(0)
|
||
|
||
print(f"Total: {total / 1e9:.2f} GB")
|
||
print(f"Allocated: {allocated / 1e9:.2f} GB")
|
||
print(f"Cached: {cached / 1e9:.2f} GB")
|
||
|
||
|
||
def inference_process(args: argparse.Namespace):
|
||
"""
|
||
Perform inference processing.
|
||
|
||
Args:
|
||
args (argparse.Namespace): Command-line arguments.
|
||
|
||
This function initializes the configuration for the inference process. It sets up the necessary
|
||
modules and variables to prepare for the upcoming inference steps.
|
||
"""
|
||
# 1. init config
|
||
cli_args = filter_non_none(vars(args))
|
||
config = OmegaConf.load(args.config)
|
||
config = OmegaConf.merge(config, cli_args)
|
||
save_path = config.save_path
|
||
if not os.path.exists(save_path):
|
||
os.makedirs(save_path)
|
||
|
||
# 2. runtime variables
|
||
device = torch.device(
|
||
"cuda") if torch.cuda.is_available() else torch.device("cpu")
|
||
if config.weight_dtype == "fp16":
|
||
weight_dtype = torch.float16
|
||
elif config.weight_dtype == "bf16":
|
||
weight_dtype = torch.bfloat16
|
||
elif config.weight_dtype == "fp32":
|
||
weight_dtype = torch.float32
|
||
else:
|
||
weight_dtype = torch.float32
|
||
|
||
# 3. prepare inference data
|
||
# 3.1 prepare source image, face mask
|
||
img_size = (config.data.source_image.width,
|
||
config.data.source_image.height)
|
||
clip_length = config.data.n_sample_frames
|
||
mask_path = config.mask_path
|
||
mask_512_path = config.mask_512_path
|
||
|
||
temp_dir = os.path.join(save_path, "img")
|
||
os.makedirs(temp_dir, exist_ok=True)
|
||
os.system(f"rm -r {temp_dir}/*")
|
||
|
||
# 3.2 prepare audio embeddings
|
||
sample_rate = config.data.driving_audio.sample_rate
|
||
assert sample_rate == 16000, "audio sample rate must be 16000"
|
||
fps = config.data.export_video.fps
|
||
|
||
### load audio processor params
|
||
audio_processor = load_audio_model(model_path=config.audio_model_path, device=device)
|
||
|
||
audio_temp_path = os.path.join(save_path, "audio.wav")
|
||
process_audio(config.driving_audio, audio_temp_path)
|
||
|
||
# 4. build modules
|
||
sched_kwargs = OmegaConf.to_container(config.noise_scheduler_kwargs)
|
||
if config.enable_zero_snr:
|
||
sched_kwargs.update(
|
||
rescale_betas_zero_snr=True,
|
||
timestep_spacing="trailing",
|
||
prediction_type="v_prediction",
|
||
)
|
||
val_noise_scheduler = DDIMScheduler(**sched_kwargs)
|
||
sched_kwargs.update({"beta_schedule": "scaled_linear"})
|
||
|
||
vae = AutoencoderKL.from_pretrained(config.vae.model_path)
|
||
reference_unet = UNet2DConditionModel.from_pretrained(
|
||
config.base_model_path, subfolder="unet")
|
||
denoising_unet = EchoUNet3DConditionModel.from_pretrained_2d(
|
||
config.base_model_path,
|
||
"",
|
||
subfolder="unet",
|
||
unet_additional_kwargs=OmegaConf.to_container(
|
||
config.unet_additional_kwargs)
|
||
)
|
||
|
||
reference_unet.load_state_dict(
|
||
torch.load(config.reference_unet_path, map_location="cpu"),
|
||
)
|
||
|
||
denoising_unet.load_state_dict(
|
||
torch.load(config.denoising_unet_path, map_location="cpu"),
|
||
)
|
||
|
||
# Freeze
|
||
vae.requires_grad_(False)
|
||
reference_unet.requires_grad_(False)
|
||
denoising_unet.requires_grad_(False)
|
||
|
||
# reference_unet.enable_gradient_checkpointing()
|
||
# denoising_unet.enable_gradient_checkpointing()
|
||
|
||
vae.eval()
|
||
reference_unet.eval()
|
||
denoising_unet.eval()
|
||
|
||
net = Net(
|
||
reference_unet,
|
||
denoising_unet
|
||
)
|
||
|
||
|
||
# 5. inference
|
||
pipeline = FaceAnimatePipeline(
|
||
vae=vae,
|
||
reference_unet=net.reference_unet,
|
||
denoising_unet=net.denoising_unet,
|
||
scheduler=val_noise_scheduler
|
||
)
|
||
|
||
|
||
pipeline.to(device=device, dtype=weight_dtype)
|
||
|
||
image_processor = ImageProcessor(img_size)
|
||
|
||
generator = torch.manual_seed(42)
|
||
|
||
video_stream = cv2.VideoCapture(args.source_video)
|
||
|
||
temp_video_path = config.output[:4] + "_no_audio.mp4"
|
||
writer = imageio.get_writer(
|
||
temp_video_path,
|
||
fps=25,
|
||
codec='libx264',
|
||
quality=10, # 0 (low) to 10 (high)
|
||
ffmpeg_params=[
|
||
'-pix_fmt', 'yuv420p',
|
||
'-crf', '18' # Lower CRF = better quality (17–23 is visually lossless range)
|
||
])
|
||
|
||
whisper_feature = audio_processor.audio2feat(audio_temp_path)
|
||
|
||
whisper_chunks = audio_processor.feature2chunks(feature_array=whisper_feature, fps=fps)
|
||
# torch.save(whisper_chunks, "./cache/audio.pt")
|
||
audio_frame_num = whisper_chunks.shape[0]
|
||
audio_fea_final = torch.Tensor(whisper_chunks).to(dtype=vae.dtype, device=vae.device)
|
||
# audio_fea_final = torch.tensor(torch.load("/workspace/preprocess/audio_emb_whisper/cM096XGOCG4_C1.pt")).to(dtype=vae.dtype, device=vae.device)
|
||
audio_fea_final = audio_fea_final.unsqueeze(0)
|
||
|
||
video_length = min(int(video_stream.get(cv2.CAP_PROP_FRAME_COUNT)), audio_frame_num)
|
||
if video_length < audio_frame_num:
|
||
audio_fea_final = audio_fea_final[:, :video_length, :, :]
|
||
|
||
times = video_length // config.inference_frames_num
|
||
|
||
frames_video = []
|
||
count_motion = 0
|
||
count_img = 0
|
||
count = 0
|
||
finish = False
|
||
|
||
while 1:
|
||
|
||
still_reading, frame = video_stream.read()
|
||
|
||
if (not still_reading) or (count_motion >= video_length):
|
||
video_stream.release()
|
||
finish = True
|
||
else:
|
||
frames_video.append(frame)
|
||
count_motion += 1
|
||
|
||
if len(frames_video)==0:
|
||
break
|
||
|
||
if ((count_motion % config.inference_frames_num) == 0) or finish:
|
||
|
||
print(f"processing {count}/{times}")
|
||
source_image_pixels_all, bbox_all, frame_list_all = image_processor.preprocess_frames_jump(frames_video)
|
||
|
||
del frames_video
|
||
|
||
for source_image_pixels_list, bbox_list, frame_list in zip(source_image_pixels_all, bbox_all, frame_list_all):
|
||
|
||
print(f"generating {count}/{times}")
|
||
|
||
if bbox_list[0]!=(0,0,0,0):
|
||
|
||
audio_tensor = audio_fea_final[:, count_img : count_img + len(frame_list)]
|
||
|
||
main_length = len(frame_list)
|
||
diff = 12 - main_length
|
||
|
||
if diff > 0:
|
||
source_image_pixels_list.extend([source_image_pixels_list[-1]]*diff)
|
||
audio_tensor = torch.cat((audio_tensor, audio_tensor[:,-1:].repeat(1,diff,1,1)), dim=1)
|
||
|
||
source_image_pixels = torch.cat(source_image_pixels_list, dim=0)
|
||
mask_main = create_mask(mask_path, mask_512_path, len(frame_list))
|
||
mask_main = mask_main.unsqueeze(0).to(device="cuda", dtype=weight_dtype)
|
||
torch.cuda.empty_cache()
|
||
|
||
with torch.no_grad():
|
||
pipeline_output = pipeline(
|
||
gt_encode=source_image_pixels,
|
||
audio_tensor=audio_tensor,
|
||
mask_main=mask_main,
|
||
width=img_size[0],
|
||
height=img_size[1],
|
||
video_length=len(source_image_pixels_list),
|
||
num_inference_steps=config.inference_steps,
|
||
guidance_scale=config.cfg_scale,
|
||
generator=generator
|
||
)
|
||
|
||
torch.cuda.empty_cache()
|
||
|
||
images_out = pipeline_output.videos.squeeze(0).permute(1, 2, 3, 0).cpu().numpy()[:main_length] # convert to [f, h, w, c]
|
||
images_out = np.clip(images_out * 255, 0, 255).astype(np.uint8) # to [0, 255]
|
||
|
||
for image_out, bbox, frame in zip(images_out, bbox_list, frame_list):
|
||
|
||
frame_h, frame_w, _ = frame.shape
|
||
path = os.path.join(temp_dir, f"image_{count_img}.png")
|
||
|
||
if bbox==(0,0,0,0):
|
||
writer.append_data(cv2.cvtColor(frame, cv2.COLOR_RGB2BGR))
|
||
# cv2.imwrite(path, frame)
|
||
else:
|
||
image_out = cv2.cvtColor(image_out, cv2.COLOR_RGB2BGR)
|
||
frame_in_out = frame.copy()
|
||
frame_in_out[bbox[1]:bbox[3], bbox[0]:bbox[2]] = cv2.resize(image_out, (bbox[2] - bbox[0], bbox[3] - bbox[1]))
|
||
mask = create_mask_alpha(frame_w, frame_h, bbox)
|
||
blurred_image = (frame * mask + frame_in_out * (1-mask)).astype(np.uint8)
|
||
writer.append_data(cv2.cvtColor(blurred_image, cv2.COLOR_RGB2BGR))
|
||
# cv2.imwrite(path, blurred_image)
|
||
|
||
count_img += 1
|
||
|
||
else:
|
||
for frame in frame_list:
|
||
path = os.path.join(temp_dir, f"image_{count_img}.png")
|
||
writer.append_data(cv2.cvtColor(frame, cv2.COLOR_RGB2BGR))
|
||
# cv2.imwrite(path, frame)
|
||
count_img += 1
|
||
|
||
count += 1
|
||
|
||
frames_video = []
|
||
|
||
if finish:
|
||
break
|
||
|
||
writer.close()
|
||
|
||
command = f"ffmpeg -i {temp_video_path} -i {config.driving_audio} -map 0:v -map 1:a -c:v h264 -shortest -y {config.output} -loglevel quiet"
|
||
subprocess.run(command, shell=True, check=True, text=True)
|
||
|
||
# output_file = config.output
|
||
# frames_path = os.path.join(temp_dir, f"image_%d.png")
|
||
# save_video(frames_path, config.driving_audio, output_file, frame_w, frame_h)
|
||
|
||
return
|
||
|
||
|
||
if __name__ == "__main__":
|
||
parser = argparse.ArgumentParser()
|
||
|
||
parser.add_argument("-c", "--config", default="configs/inference/default.yaml")
|
||
parser.add_argument("--source_video", type=str, required=False, help="source image")
|
||
parser.add_argument("--driving_audio", type=str, required=False, help="driving audio")
|
||
parser.add_argument("--output", type=str, help="output video file name", default="./cache/output.mp4")
|
||
parser.add_argument("--inference_frames_num", type=int, help="inference frames number", default=1500)
|
||
|
||
|
||
command_line_args = parser.parse_args()
|
||
|
||
inference_process(command_line_args)
|