282 lines
10 KiB
Python
282 lines
10 KiB
Python
# pylint: disable=E1101
|
|
# scripts/inference.py
|
|
|
|
"""
|
|
This script contains the main inference pipeline for processing audio and image inputs to generate a video output.
|
|
|
|
The script imports necessary packages and classes, defines a neural network model,
|
|
and contains functions for processing audio embeddings and performing inference.
|
|
|
|
The main inference process is outlined in the following steps:
|
|
1. Initialize the configuration.
|
|
2. Set up runtime variables.
|
|
3. Prepare the input data for inference (source image, face mask, and face embeddings).
|
|
4. Process the audio embeddings.
|
|
5. Build and freeze the model and scheduler.
|
|
6. Run the inference loop and save the result.
|
|
|
|
Usage:
|
|
This script can be run from the command line with the following arguments:
|
|
- audio_path: Path to the audio file.
|
|
- image_path: Path to the source image.
|
|
- face_mask_path: Path to the face mask image.
|
|
- face_emb_path: Path to the face embeddings file.
|
|
- output_path: Path to save the output video.
|
|
|
|
Example:
|
|
python scripts/inference.py --audio_path audio.wav --image_path image.jpg
|
|
--face_mask_path face_mask.png --face_emb_path face_emb.pt --output_path output.mp4
|
|
"""
|
|
|
|
import argparse
|
|
import os
|
|
import sys
|
|
|
|
import torch
|
|
from diffusers import AutoencoderKL, DDIMScheduler
|
|
from omegaconf import OmegaConf
|
|
from torch import nn
|
|
from pathlib import Path
|
|
import numpy as np
|
|
import torchvision.transforms as transforms
|
|
from PIL import Image
|
|
from pydub import AudioSegment
|
|
|
|
sys.path.append(os.path.dirname(os.path.dirname(os.path.abspath(__file__))))
|
|
|
|
from ..hallo.animate.face_animate import FaceAnimatePipeline
|
|
from ..hallo.datasets.audio_processor import AudioProcessor
|
|
from ..hallo.datasets.image_processor import ImageProcessor
|
|
from ..hallo.models.audio_proj import AudioProjModel
|
|
from ..hallo.models.face_locator import FaceLocator
|
|
from ..hallo.models.image_proj import ImageProjModel
|
|
from ..hallo.models.unet_2d_condition import UNet2DConditionModel
|
|
from ..hallo.models.unet_3d import UNet3DConditionModel
|
|
from ..hallo.utils.config import filter_non_none
|
|
from ..hallo.utils.util import tensor_to_video_batch, merge_videos
|
|
import folder_paths
|
|
|
|
from icecream import ic
|
|
|
|
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.
|
|
face_locator (FaceLocator): The FaceLocator model used to locate the face in the input image.
|
|
imageproj (nn.Module): The ImageProjector model used to project the source image onto the face.
|
|
audioproj (nn.Module): The AudioProjector model used to project the audio embeddings onto the face.
|
|
"""
|
|
def __init__(
|
|
self,
|
|
reference_unet: UNet2DConditionModel,
|
|
denoising_unet: UNet3DConditionModel,
|
|
face_locator: FaceLocator,
|
|
imageproj,
|
|
audioproj,
|
|
):
|
|
super().__init__()
|
|
self.reference_unet = reference_unet
|
|
self.denoising_unet = denoising_unet
|
|
self.face_locator = face_locator
|
|
self.imageproj = imageproj
|
|
self.audioproj = audioproj
|
|
|
|
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,
|
|
"face_locator": self.face_locator,
|
|
"imageproj": self.imageproj,
|
|
"audioproj": self.audioproj,
|
|
}
|
|
|
|
|
|
def process_audio_emb(audio_emb):
|
|
"""
|
|
Process the audio embedding to concatenate with other tensors.
|
|
|
|
Parameters:
|
|
audio_emb (torch.Tensor): The audio embedding tensor to process.
|
|
|
|
Returns:
|
|
concatenated_tensors (List[torch.Tensor]): The concatenated tensor list.
|
|
"""
|
|
concatenated_tensors = []
|
|
|
|
for i in range(audio_emb.shape[0]):
|
|
vectors_to_concat = [
|
|
audio_emb[max(min(i + j, audio_emb.shape[0]-1), 0)]for j in range(-2, 3)]
|
|
concatenated_tensors.append(torch.stack(vectors_to_concat, dim=0))
|
|
|
|
audio_emb = torch.stack(concatenated_tensors, dim=0)
|
|
|
|
return audio_emb
|
|
|
|
def save_image_batch(image_tensor, save_path):
|
|
image_tensor = (image_tensor + 1) / 2
|
|
|
|
os.makedirs(save_path, exist_ok=True)
|
|
|
|
for i in range(image_tensor.shape[0]):
|
|
img_tensor = image_tensor[i]
|
|
|
|
img_array = img_tensor.permute(1, 2, 0).cpu().numpy()
|
|
|
|
img_array = (img_array * 255).astype(np.uint8)
|
|
|
|
image = Image.fromarray(img_array)
|
|
image.save(os.path.join(save_path, f'motion_frame_{i}.png'))
|
|
|
|
|
|
def cut_audio(audio_path, save_dir, length=60):
|
|
audio = AudioSegment.from_wav(audio_path)
|
|
|
|
segment_length = length * 1000 # pydub使用毫秒
|
|
|
|
num_segments = len(audio) // segment_length + (1 if len(audio) % segment_length != 0 else 0)
|
|
|
|
os.makedirs(save_dir, exist_ok=True)
|
|
|
|
audio_list = []
|
|
|
|
for i in range(num_segments):
|
|
start_time = i * segment_length
|
|
end_time = min((i + 1) * segment_length, len(audio))
|
|
segment = audio[start_time:end_time]
|
|
|
|
path = f"{save_dir}/segment_{i+1}.wav"
|
|
audio_list.append(path)
|
|
segment.export(path, format="wav")
|
|
|
|
return audio_list
|
|
|
|
|
|
def inference_process(driving_audio_path,pose_weight,face_weight,lip_weight,audio_emb,clip_length,audio_length,source_image_pixels,source_image_face_region,
|
|
source_image_face_emb,source_image_full_mask,source_image_face_mask,source_image_lip_mask,img_size,net,pipeline,inference_steps,cfg_scale,seed,use_mask,save_video,fps):
|
|
|
|
|
|
motion_scale = [pose_weight, face_weight, lip_weight] # 1.0,1.0,1.0
|
|
# from here
|
|
times = audio_emb.shape[0] // clip_length
|
|
tensor_result = []
|
|
generator = torch.manual_seed(seed)
|
|
|
|
ic(audio_emb.shape)
|
|
ic(audio_length)
|
|
batch_size = 60
|
|
start = 0
|
|
for t in range(times):
|
|
print(f"Start infer batch {t+1} ,and total batchs is {times}.")
|
|
|
|
if len(tensor_result) == 0:
|
|
# The first iteration
|
|
motion_zeros = source_image_pixels.repeat(
|
|
2, 1, 1, 1)
|
|
motion_zeros = motion_zeros.to(
|
|
dtype=source_image_pixels.dtype, device=source_image_pixels.device)
|
|
pixel_values_ref_img = torch.cat(
|
|
[source_image_pixels, motion_zeros], dim=0) # concat the ref image and the first motion frames
|
|
else:
|
|
motion_frames = tensor_result[-1][0]
|
|
motion_frames = motion_frames.permute(1, 0, 2, 3)
|
|
motion_frames = motion_frames[0-2:]
|
|
motion_frames = motion_frames * 2.0 - 1.0
|
|
motion_frames = motion_frames.to(
|
|
dtype=source_image_pixels.dtype, device=source_image_pixels.device)
|
|
pixel_values_ref_img = torch.cat(
|
|
[source_image_pixels, motion_frames], dim=0) # concat the ref image and the motion frames
|
|
|
|
pixel_values_ref_img = pixel_values_ref_img.unsqueeze(0)
|
|
|
|
pixel_motion_values = pixel_values_ref_img[:, 1:]
|
|
|
|
if use_mask:
|
|
b, f, c, h, w = pixel_motion_values.shape
|
|
rand_mask = torch.rand(h, w)
|
|
mask = rand_mask > 0.25
|
|
mask = mask.unsqueeze(0).unsqueeze(0).unsqueeze(0)
|
|
mask = mask.expand(b, f, c, h, w)
|
|
|
|
face_mask = source_image_face_region.repeat(f, 1, 1, 1).unsqueeze(0)
|
|
assert face_mask.shape == mask.shape
|
|
mask = mask | face_mask.bool()
|
|
|
|
pixel_motion_values = pixel_motion_values * mask
|
|
pixel_values_ref_img[:, 1:] = pixel_motion_values
|
|
|
|
|
|
assert pixel_motion_values.shape[0] == 1
|
|
|
|
audio_tensor = audio_emb[
|
|
t * clip_length: min((t + 1) * clip_length, audio_emb.shape[0])
|
|
]
|
|
audio_tensor = audio_tensor.unsqueeze(0)
|
|
audio_tensor = audio_tensor.to(
|
|
device=net.audioproj.device, dtype=net.audioproj.dtype)
|
|
audio_tensor = net.audioproj(audio_tensor)
|
|
|
|
pipeline_output = pipeline(
|
|
ref_image=pixel_values_ref_img,
|
|
audio_tensor=audio_tensor,
|
|
face_emb=source_image_face_emb,
|
|
face_mask=source_image_face_region,
|
|
pixel_values_full_mask=source_image_full_mask,
|
|
pixel_values_face_mask=source_image_face_mask,
|
|
pixel_values_lip_mask=source_image_lip_mask,
|
|
width=img_size[0],
|
|
height=img_size[1],
|
|
video_length=clip_length,
|
|
num_inference_steps=inference_steps,
|
|
guidance_scale=cfg_scale,
|
|
generator=generator,
|
|
motion_scale=motion_scale,
|
|
)
|
|
|
|
ic(pipeline_output.videos.shape)# [1,3, 16, 512, 512]
|
|
tensor_result.append(pipeline_output.videos)
|
|
|
|
if (t+1) % batch_size == 0 or (t+1)==times:
|
|
last_motion_frame = [tensor_result[-1]]
|
|
ic(len(tensor_result))
|
|
|
|
if start!=0:
|
|
tensor_result = torch.cat(tensor_result[1:], dim=2)
|
|
else:
|
|
tensor_result = torch.cat(tensor_result, dim=2)
|
|
|
|
tensor_result = tensor_result.squeeze(0)
|
|
f = tensor_result.shape[1]
|
|
length = min(f, audio_length)
|
|
tensor_result = tensor_result[:, :length] #torch.Size([3, 188, 512, 512])
|
|
tensor_list= tensor_result.permute(1, 2, 3, 0) # convert to [f, h, w, c] comfyUI [B,H,W,C], C=3
|
|
|
|
ic(tensor_result.shape)
|
|
ic(start)
|
|
ic(audio_length)
|
|
|
|
output_file = os.path.join(folder_paths.get_output_directory(), f"Hallo2-{t+1:06}.mp4")
|
|
if save_video:
|
|
tensor_to_video_batch(tensor_result, output_file, start, driving_audio_path, fps)
|
|
output_path = output_file
|
|
del tensor_result
|
|
tensor_result = last_motion_frame
|
|
audio_length -= length
|
|
start += length
|
|
else:
|
|
output_path=""
|
|
torch.cuda.empty_cache()
|
|
return tensor_list,output_path
|
|
|
|
|