Support for Hunyuan Video Avatar

This commit is contained in:
DeepBeepMeep
2025-06-05 15:35:58 +02:00
parent 7670af9610
commit 1976868f6a
15 changed files with 2562 additions and 81 deletions

View File

@@ -0,0 +1,170 @@
import os
import cv2
import math
import json
import torch
import random
import librosa
import traceback
import torchvision
import numpy as np
import pandas as pd
from PIL import Image
from einops import rearrange
from torch.utils.data import Dataset
from decord import VideoReader, cpu
from transformers import CLIPImageProcessor
import torchvision.transforms as transforms
from torchvision.transforms import ToPILImage
def get_audio_feature(feature_extractor, audio_path):
audio_input, sampling_rate = librosa.load(audio_path, sr=16000)
assert sampling_rate == 16000
audio_features = []
window = 750*640
for i in range(0, len(audio_input), window):
audio_feature = feature_extractor(audio_input[i:i+window],
sampling_rate=sampling_rate,
return_tensors="pt",
).input_features
audio_features.append(audio_feature)
audio_features = torch.cat(audio_features, dim=-1)
return audio_features, len(audio_input) // 640
class VideoAudioTextLoaderVal(Dataset):
def __init__(
self,
image_size: int,
meta_file: str,
**kwargs,
):
super().__init__()
self.meta_file = meta_file
self.image_size = image_size
self.text_encoder = kwargs.get("text_encoder", None) # llava_text_encoder
self.text_encoder_2 = kwargs.get("text_encoder_2", None) # clipL_text_encoder
self.feature_extractor = kwargs.get("feature_extractor", None)
self.meta_files = []
csv_data = pd.read_csv(meta_file)
for idx in range(len(csv_data)):
self.meta_files.append(
{
"videoid": str(csv_data["videoid"][idx]),
"image_path": str(csv_data["image"][idx]),
"audio_path": str(csv_data["audio"][idx]),
"prompt": str(csv_data["prompt"][idx]),
"fps": float(csv_data["fps"][idx])
}
)
self.llava_transform = transforms.Compose(
[
transforms.Resize((336, 336), interpolation=transforms.InterpolationMode.BILINEAR),
transforms.ToTensor(),
transforms.Normalize((0.48145466, 0.4578275, 0.4082107), (0.26862954, 0.26130258, 0.27577711)),
]
)
self.clip_image_processor = CLIPImageProcessor()
self.device = torch.device("cuda")
self.weight_dtype = torch.float16
def __len__(self):
return len(self.meta_files)
@staticmethod
def get_text_tokens(text_encoder, description, dtype_encode="video"):
text_inputs = text_encoder.text2tokens(description, data_type=dtype_encode)
text_ids = text_inputs["input_ids"].squeeze(0)
text_mask = text_inputs["attention_mask"].squeeze(0)
return text_ids, text_mask
def get_batch_data(self, idx):
meta_file = self.meta_files[idx]
videoid = meta_file["videoid"]
image_path = meta_file["image_path"]
audio_path = meta_file["audio_path"]
prompt = "Authentic, Realistic, Natural, High-quality, Lens-Fixed, " + meta_file["prompt"]
fps = meta_file["fps"]
img_size = self.image_size
ref_image = Image.open(image_path).convert('RGB')
# Resize reference image
w, h = ref_image.size
scale = img_size / min(w, h)
new_w = round(w * scale / 64) * 64
new_h = round(h * scale / 64) * 64
if img_size == 704:
img_size_long = 1216
if new_w * new_h > img_size * img_size_long:
import math
scale = math.sqrt(img_size * img_size_long / w / h)
new_w = round(w * scale / 64) * 64
new_h = round(h * scale / 64) * 64
ref_image = ref_image.resize((new_w, new_h), Image.LANCZOS)
ref_image = np.array(ref_image)
ref_image = torch.from_numpy(ref_image)
audio_input, audio_len = get_audio_feature(self.feature_extractor, audio_path)
audio_prompts = audio_input[0]
motion_bucket_id_heads = np.array([25] * 4)
motion_bucket_id_exps = np.array([30] * 4)
motion_bucket_id_heads = torch.from_numpy(motion_bucket_id_heads)
motion_bucket_id_exps = torch.from_numpy(motion_bucket_id_exps)
fps = torch.from_numpy(np.array(fps))
to_pil = ToPILImage()
pixel_value_ref = rearrange(ref_image.clone().unsqueeze(0), "b h w c -> b c h w") # (b c h w)
pixel_value_ref_llava = [self.llava_transform(to_pil(image)) for image in pixel_value_ref]
pixel_value_ref_llava = torch.stack(pixel_value_ref_llava, dim=0)
pixel_value_ref_clip = self.clip_image_processor(
images=Image.fromarray((pixel_value_ref[0].permute(1,2,0)).data.cpu().numpy().astype(np.uint8)),
return_tensors="pt"
).pixel_values[0]
pixel_value_ref_clip = pixel_value_ref_clip.unsqueeze(0)
# Encode text prompts
text_ids, text_mask = self.get_text_tokens(self.text_encoder, prompt)
text_ids_2, text_mask_2 = self.get_text_tokens(self.text_encoder_2, prompt)
# Output batch
batch = {
"text_prompt": prompt, #
"videoid": videoid,
"pixel_value_ref": pixel_value_ref.to(dtype=torch.float16), # 参考图用于vae提特征 (1, 3, h, w), 取值范围(0, 255)
"pixel_value_ref_llava": pixel_value_ref_llava.to(dtype=torch.float16), # 参考图用于llava提特征 (1, 3, 336, 336), 取值范围 = CLIP取值范围
"pixel_value_ref_clip": pixel_value_ref_clip.to(dtype=torch.float16), # 参考图用于clip_image_encoder提特征 (1, 3, 244, 244), 取值范围 = CLIP取值范围
"audio_prompts": audio_prompts.to(dtype=torch.float16),
"motion_bucket_id_heads": motion_bucket_id_heads.to(dtype=text_ids.dtype),
"motion_bucket_id_exps": motion_bucket_id_exps.to(dtype=text_ids.dtype),
"fps": fps.to(dtype=torch.float16),
"text_ids": text_ids.clone(), # 对应llava_text_encoder
"text_mask": text_mask.clone(), # 对应llava_text_encoder
"text_ids_2": text_ids_2.clone(), # 对应clip_text_encoder
"text_mask_2": text_mask_2.clone(), # 对应clip_text_encoder
"audio_len": audio_len,
"image_path": image_path,
"audio_path": audio_path,
}
return batch
def __getitem__(self, idx):
return self.get_batch_data(idx)

View File

@@ -0,0 +1,72 @@
import os
import cv2
import json
import time
import decord
import einops
import librosa
import torch
import random
import argparse
import traceback
import numpy as np
from tqdm import tqdm
from PIL import Image
from einops import rearrange
def get_facemask(ref_image, align_instance, area=1.25):
# ref_image: (b f c h w)
bsz, f, c, h, w = ref_image.shape
images = rearrange(ref_image, "b f c h w -> (b f) h w c").data.cpu().numpy().astype(np.uint8)
face_masks = []
for image in images:
image_pil = Image.fromarray(image).convert("RGB")
_, _, bboxes_list = align_instance(np.array(image_pil)[:,:,[2,1,0]], maxface=True)
try:
bboxSrc = bboxes_list[0]
except:
bboxSrc = [0, 0, w, h]
x1, y1, ww, hh = bboxSrc
x2, y2 = x1 + ww, y1 + hh
ww, hh = (x2-x1) * area, (y2-y1) * area
center = [(x2+x1)//2, (y2+y1)//2]
x1 = max(center[0] - ww//2, 0)
y1 = max(center[1] - hh//2, 0)
x2 = min(center[0] + ww//2, w)
y2 = min(center[1] + hh//2, h)
face_mask = np.zeros_like(np.array(image_pil))
face_mask[int(y1):int(y2), int(x1):int(x2)] = 1.0
face_masks.append(torch.from_numpy(face_mask[...,:1]))
face_masks = torch.stack(face_masks, dim=0) # (b*f, h, w, c)
face_masks = rearrange(face_masks, "(b f) h w c -> b c f h w", b=bsz, f=f)
face_masks = face_masks.to(device=ref_image.device, dtype=ref_image.dtype)
return face_masks
def encode_audio(wav2vec, audio_feats, fps, num_frames=129):
if fps == 25:
start_ts = [0]
step_ts = [1]
elif fps == 12.5:
start_ts = [0]
step_ts = [2]
num_frames = min(num_frames, 400)
audio_feats = wav2vec.encoder(audio_feats.unsqueeze(0)[:, :, :3000], output_hidden_states=True).hidden_states
audio_feats = torch.stack(audio_feats, dim=2)
audio_feats = torch.cat([torch.zeros_like(audio_feats[:,:4]), audio_feats], 1)
audio_prompts = []
for bb in range(1):
audio_feats_list = []
for f in range(num_frames):
cur_t = (start_ts[bb] + f * step_ts[bb]) * 2
audio_clip = audio_feats[bb:bb+1, cur_t: cur_t+10]
audio_feats_list.append(audio_clip)
audio_feats_list = torch.stack(audio_feats_list, 1)
audio_prompts.append(audio_feats_list)
audio_prompts = torch.cat(audio_prompts)
return audio_prompts

View File

@@ -0,0 +1,41 @@
import os
import cv2
import torch
import numpy as np
import imageio
import torchvision
from einops import rearrange
def save_videos_grid(videos: torch.Tensor, path: str, rescale=False, n_rows=6, fps=8, quality=8):
videos = rearrange(videos, "b c t h w -> t b c h w")
outputs = []
for x in videos:
x = torchvision.utils.make_grid(x, nrow=n_rows)
x = x.transpose(0, 1).transpose(1, 2).squeeze(-1)
if rescale:
x = (x + 1.0) / 2.0 # -1,1 -> 0,1
x = torch.clamp(x,0,1)
x = (x * 255).numpy().astype(np.uint8)
outputs.append(x)
os.makedirs(os.path.dirname(path), exist_ok=True)
imageio.mimsave(path, outputs, fps=fps, quality=quality)
def pad_image(crop_img, size, color=(255, 255, 255), resize_ratio=1):
crop_h, crop_w = crop_img.shape[:2]
target_w, target_h = size
scale_h, scale_w = target_h / crop_h, target_w / crop_w
if scale_w > scale_h:
resize_h = int(target_h*resize_ratio)
resize_w = int(crop_w / crop_h * resize_h)
else:
resize_w = int(target_w*resize_ratio)
resize_h = int(crop_h / crop_w * resize_w)
crop_img = cv2.resize(crop_img, (resize_w, resize_h))
pad_left = (target_w - resize_w) // 2
pad_top = (target_h - resize_h) // 2
pad_right = target_w - resize_w - pad_left
pad_bottom = target_h - resize_h - pad_top
crop_img = cv2.copyMakeBorder(crop_img, pad_top, pad_bottom, pad_left, pad_right, cv2.BORDER_CONSTANT, value=color)
return crop_img

View File

@@ -0,0 +1 @@
from .align import AlignImage

View File

@@ -0,0 +1,34 @@
import os
import sys
import torch
from .detface import DetFace
class AlignImage(object):
def __init__(self, device='cuda', det_path=''):
self.facedet = DetFace(pt_path=det_path, confThreshold=0.5, nmsThreshold=0.45, device=device)
@torch.no_grad()
def __call__(self, im, maxface=False):
bboxes, kpss, scores = self.facedet.detect(im)
face_num = bboxes.shape[0]
five_pts_list = []
scores_list = []
bboxes_list = []
for i in range(face_num):
five_pts_list.append(kpss[i].reshape(5,2))
scores_list.append(scores[i])
bboxes_list.append(bboxes[i])
if maxface and face_num>1:
max_idx = 0
max_area = (bboxes[0, 2])*(bboxes[0, 3])
for i in range(1, face_num):
area = (bboxes[i,2])*(bboxes[i,3])
if area>max_area:
max_idx = i
five_pts_list = [five_pts_list[max_idx]]
scores_list = [scores_list[max_idx]]
bboxes_list = [bboxes_list[max_idx]]
return five_pts_list, scores_list, bboxes_list

View File

@@ -0,0 +1,283 @@
# -*- coding: UTF-8 -*-
import os
import cv2
import numpy as np
import torch
import torchvision
def xyxy2xywh(x):
# Convert nx4 boxes from [x1, y1, x2, y2] to [x, y, w, h] where xy1=top-left, xy2=bottom-right
y = x.clone() if isinstance(x, torch.Tensor) else np.copy(x)
y[:, 0] = (x[:, 0] + x[:, 2]) / 2 # x center
y[:, 1] = (x[:, 1] + x[:, 3]) / 2 # y center
y[:, 2] = x[:, 2] - x[:, 0] # width
y[:, 3] = x[:, 3] - x[:, 1] # height
return y
def xywh2xyxy(x):
# Convert nx4 boxes from [x, y, w, h] to [x1, y1, x2, y2] where xy1=top-left, xy2=bottom-right
y = x.clone() if isinstance(x, torch.Tensor) else np.copy(x)
y[:, 0] = x[:, 0] - x[:, 2] / 2 # top left x
y[:, 1] = x[:, 1] - x[:, 3] / 2 # top left y
y[:, 2] = x[:, 0] + x[:, 2] / 2 # bottom right x
y[:, 3] = x[:, 1] + x[:, 3] / 2 # bottom right y
return y
def box_iou(box1, box2):
# https://github.com/pytorch/vision/blob/master/torchvision/ops/boxes.py
"""
Return intersection-over-union (Jaccard index) of boxes.
Both sets of boxes are expected to be in (x1, y1, x2, y2) format.
Arguments:
box1 (Tensor[N, 4])
box2 (Tensor[M, 4])
Returns:
iou (Tensor[N, M]): the NxM matrix containing the pairwise
IoU values for every element in boxes1 and boxes2
"""
def box_area(box):
# box = 4xn
return (box[2] - box[0]) * (box[3] - box[1])
area1 = box_area(box1.T)
area2 = box_area(box2.T)
# inter(N,M) = (rb(N,M,2) - lt(N,M,2)).clamp(0).prod(2)
inter = (torch.min(box1[:, None, 2:], box2[:, 2:]) -
torch.max(box1[:, None, :2], box2[:, :2])).clamp(0).prod(2)
# iou = inter / (area1 + area2 - inter)
return inter / (area1[:, None] + area2 - inter)
def scale_coords(img1_shape, coords, img0_shape, ratio_pad=None):
# Rescale coords (xyxy) from img1_shape to img0_shape
if ratio_pad is None: # calculate from img0_shape
gain = min(img1_shape[0] / img0_shape[0], img1_shape[1] / img0_shape[1]) # gain = old / new
pad = (img1_shape[1] - img0_shape[1] * gain) / 2, (img1_shape[0] - img0_shape[0] * gain) / 2 # wh padding
else:
gain = ratio_pad[0][0]
pad = ratio_pad[1]
coords[:, [0, 2]] -= pad[0] # x padding
coords[:, [1, 3]] -= pad[1] # y padding
coords[:, :4] /= gain
clip_coords(coords, img0_shape)
return coords
def clip_coords(boxes, img_shape):
# Clip bounding xyxy bounding boxes to image shape (height, width)
boxes[:, 0].clamp_(0, img_shape[1]) # x1
boxes[:, 1].clamp_(0, img_shape[0]) # y1
boxes[:, 2].clamp_(0, img_shape[1]) # x2
boxes[:, 3].clamp_(0, img_shape[0]) # y2
def scale_coords_landmarks(img1_shape, coords, img0_shape, ratio_pad=None):
# Rescale coords (xyxy) from img1_shape to img0_shape
if ratio_pad is None: # calculate from img0_shape
gain = min(img1_shape[0] / img0_shape[0], img1_shape[1] / img0_shape[1]) # gain = old / new
pad = (img1_shape[1] - img0_shape[1] * gain) / 2, (img1_shape[0] - img0_shape[0] * gain) / 2 # wh padding
else:
gain = ratio_pad[0][0]
pad = ratio_pad[1]
coords[:, [0, 2, 4, 6, 8]] -= pad[0] # x padding
coords[:, [1, 3, 5, 7, 9]] -= pad[1] # y padding
coords[:, :10] /= gain
#clip_coords(coords, img0_shape)
coords[:, 0].clamp_(0, img0_shape[1]) # x1
coords[:, 1].clamp_(0, img0_shape[0]) # y1
coords[:, 2].clamp_(0, img0_shape[1]) # x2
coords[:, 3].clamp_(0, img0_shape[0]) # y2
coords[:, 4].clamp_(0, img0_shape[1]) # x3
coords[:, 5].clamp_(0, img0_shape[0]) # y3
coords[:, 6].clamp_(0, img0_shape[1]) # x4
coords[:, 7].clamp_(0, img0_shape[0]) # y4
coords[:, 8].clamp_(0, img0_shape[1]) # x5
coords[:, 9].clamp_(0, img0_shape[0]) # y5
return coords
def show_results(img, xywh, conf, landmarks, class_num):
h,w,c = img.shape
tl = 1 or round(0.002 * (h + w) / 2) + 1 # line/font thickness
x1 = int(xywh[0] * w - 0.5 * xywh[2] * w)
y1 = int(xywh[1] * h - 0.5 * xywh[3] * h)
x2 = int(xywh[0] * w + 0.5 * xywh[2] * w)
y2 = int(xywh[1] * h + 0.5 * xywh[3] * h)
cv2.rectangle(img, (x1,y1), (x2, y2), (0,255,0), thickness=tl, lineType=cv2.LINE_AA)
clors = [(255,0,0),(0,255,0),(0,0,255),(255,255,0),(0,255,255)]
for i in range(5):
point_x = int(landmarks[2 * i] * w)
point_y = int(landmarks[2 * i + 1] * h)
cv2.circle(img, (point_x, point_y), tl+1, clors[i], -1)
tf = max(tl - 1, 1) # font thickness
label = str(conf)[:5]
cv2.putText(img, label, (x1, y1 - 2), 0, tl / 3, [225, 255, 255], thickness=tf, lineType=cv2.LINE_AA)
return img
def make_divisible(x, divisor):
# Returns x evenly divisible by divisor
return (x // divisor) * divisor
def non_max_suppression_face(prediction, conf_thres=0.5, iou_thres=0.45, classes=None, agnostic=False, labels=()):
"""Performs Non-Maximum Suppression (NMS) on inference results
Returns:
detections with shape: nx6 (x1, y1, x2, y2, conf, cls)
"""
nc = prediction.shape[2] - 15 # number of classes
xc = prediction[..., 4] > conf_thres # candidates
# Settings
min_wh, max_wh = 2, 4096 # (pixels) minimum and maximum box width and height
# time_limit = 10.0 # seconds to quit after
redundant = True # require redundant detections
multi_label = nc > 1 # multiple labels per box (adds 0.5ms/img)
merge = False # use merge-NMS
# t = time.time()
output = [torch.zeros((0, 16), device=prediction.device)] * prediction.shape[0]
for xi, x in enumerate(prediction): # image index, image inference
# Apply constraints
# x[((x[..., 2:4] < min_wh) | (x[..., 2:4] > max_wh)).any(1), 4] = 0 # width-height
x = x[xc[xi]] # confidence
# Cat apriori labels if autolabelling
if labels and len(labels[xi]):
l = labels[xi]
v = torch.zeros((len(l), nc + 15), device=x.device)
v[:, :4] = l[:, 1:5] # box
v[:, 4] = 1.0 # conf
v[range(len(l)), l[:, 0].long() + 15] = 1.0 # cls
x = torch.cat((x, v), 0)
# If none remain process next image
if not x.shape[0]:
continue
# Compute conf
x[:, 15:] *= x[:, 4:5] # conf = obj_conf * cls_conf
# Box (center x, center y, width, height) to (x1, y1, x2, y2)
box = xywh2xyxy(x[:, :4])
# Detections matrix nx6 (xyxy, conf, landmarks, cls)
if multi_label:
i, j = (x[:, 15:] > conf_thres).nonzero(as_tuple=False).T
x = torch.cat((box[i], x[i, j + 15, None], x[i, 5:15] ,j[:, None].float()), 1)
else: # best class only
conf, j = x[:, 15:].max(1, keepdim=True)
x = torch.cat((box, conf, x[:, 5:15], j.float()), 1)[conf.view(-1) > conf_thres]
# Filter by class
if classes is not None:
x = x[(x[:, 5:6] == torch.tensor(classes, device=x.device)).any(1)]
# If none remain process next image
n = x.shape[0] # number of boxes
if not n:
continue
# Batched NMS
c = x[:, 15:16] * (0 if agnostic else max_wh) # classes
boxes, scores = x[:, :4] + c, x[:, 4] # boxes (offset by class), scores
i = torchvision.ops.nms(boxes, scores, iou_thres) # NMS
#if i.shape[0] > max_det: # limit detections
# i = i[:max_det]
if merge and (1 < n < 3E3): # Merge NMS (boxes merged using weighted mean)
# update boxes as boxes(i,4) = weights(i,n) * boxes(n,4)
iou = box_iou(boxes[i], boxes) > iou_thres # iou matrix
weights = iou * scores[None] # box weights
x[i, :4] = torch.mm(weights, x[:, :4]).float() / weights.sum(1, keepdim=True) # merged boxes
if redundant:
i = i[iou.sum(1) > 1] # require redundancy
output[xi] = x[i]
# if (time.time() - t) > time_limit:
# break # time limit exceeded
return output
class DetFace():
def __init__(self, pt_path, confThreshold=0.5, nmsThreshold=0.45, device='cuda'):
assert os.path.exists(pt_path)
self.inpSize = 416
self.conf_thres = confThreshold
self.iou_thres = nmsThreshold
self.test_device = torch.device(device if torch.cuda.is_available() else "cpu")
self.model = torch.jit.load(pt_path).to(self.test_device)
self.last_w = 416
self.last_h = 416
self.grids = None
@torch.no_grad()
def detect(self, srcimg):
# t0=time.time()
h0, w0 = srcimg.shape[:2] # orig hw
r = self.inpSize / min(h0, w0) # resize image to img_size
h1 = int(h0*r+31)//32*32
w1 = int(w0*r+31)//32*32
img = cv2.resize(srcimg, (w1,h1), interpolation=cv2.INTER_LINEAR)
# Convert
img = cv2.cvtColor(img, cv2.COLOR_BGR2RGB) # BGR to RGB
# Run inference
img = torch.from_numpy(img).to(self.test_device).permute(2,0,1)
img = img.float()/255 # uint8 to fp16/32 0-1
if img.ndimension() == 3:
img = img.unsqueeze(0)
# Inference
if h1 != self.last_h or w1 != self.last_w or self.grids is None:
grids = []
for scale in [8,16,32]:
ny = h1//scale
nx = w1//scale
yv, xv = torch.meshgrid([torch.arange(ny), torch.arange(nx)])
grid = torch.stack((xv, yv), 2).view((1,1,ny, nx, 2)).float()
grids.append(grid.to(self.test_device))
self.grids = grids
self.last_w = w1
self.last_h = h1
pred = self.model(img, self.grids).cpu()
# Apply NMS
det = non_max_suppression_face(pred, self.conf_thres, self.iou_thres)[0]
# Process detections
# det = pred[0]
bboxes = np.zeros((det.shape[0], 4))
kpss = np.zeros((det.shape[0], 5, 2))
scores = np.zeros((det.shape[0]))
# gn = torch.tensor([w0, h0, w0, h0]).to(pred) # normalization gain whwh
# gn_lks = torch.tensor([w0, h0, w0, h0, w0, h0, w0, h0, w0, h0]).to(pred) # normalization gain landmarks
det = det.cpu().numpy()
for j in range(det.shape[0]):
# xywh = (xyxy2xywh(det[j, :4].view(1, 4)) / gn).view(4).cpu().numpy()
bboxes[j, 0] = det[j, 0] * w0/w1
bboxes[j, 1] = det[j, 1] * h0/h1
bboxes[j, 2] = det[j, 2] * w0/w1 - bboxes[j, 0]
bboxes[j, 3] = det[j, 3] * h0/h1 - bboxes[j, 1]
scores[j] = det[j, 4]
# landmarks = (det[j, 5:15].view(1, 10) / gn_lks).view(5,2).cpu().numpy()
kpss[j, :, :] = det[j, 5:15].reshape(5, 2) * np.array([[w0/w1,h0/h1]])
# class_num = det[j, 15].cpu().numpy()
# orgimg = show_results(orgimg, xywh, conf, landmarks, class_num)
return bboxes, kpss, scores