ИИ-Дед Мороз: создаём новогодние видео-открытки с YandexART и YandexGPT
Работа с моделью EchoMimic
Создадим Python‑класс обёртку для модели EchoMimic:
import os
import random
from datetime import datetime
from pathlib import Path
import cv2
import numpy as np
import torch
from diffusers import AutoencoderKL, DDIMScheduler
from omegaconf import OmegaConf
from PIL import Image
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 src.pipelines.pipeline_echo_mimic_acc import Audio2VideoPipeline
from src.utils.util import save_videos_grid, crop_and_pad
from src.models.face_locator import FaceLocator
from moviepy.editor import VideoFileClip, AudioFileClip
from facenet_pytorch import MTCNN
"""
Класс для преобразования аудио в видео с учетом параметров лица и движений.
"""
class Audio2Video:
def __init__(self, config="./configs/prompts/animation_acc.yaml", device="cuda", seed=420, W=512, H=512, L=1200, fps=24, sample_rate=16000, context_frames=12, context_overlap=3, facemusk_dilation_ratio=0.1, facecrop_dilation_ratio=0.5, cfg=1.0, steps=6):
"""
Инициализация класса Audio2Video.
:param config: Путь до конфигурационного файла.
:param device: Устройство для выполнения вычислений (например, "cuda" или "cpu").
:param seed: Значение для начальной инициализации генератора случайных чисел.
:param W: Ширина выходного видео.
:param H: Высота выходного видео.
:param L: Длительность видео в кадрах.
:param fps: Количество кадров в секунду.
:param sample_rate: Частота дискретизации аудио.
:param context_frames: Количество кадров контекста для обработки.
:param context_overlap: Перекрытие между контекстными кадрами.
:param facemusk_dilation_ratio: Доля увеличения маски лица для плавности переходов.
:param facecrop_dilation_ratio: Доля увеличения области обрезки лица.
:param cfg: Вес конфигурации для управления.
:param steps: Количество шагов генерации.
"""
# Инициализация параметров
self.config = config
self.device = device
self.seed = seed
self.W = W
self.H = H
self.L = L
self.fps = fps
self.sample_rate = sample_rate
self.context_frames = context_frames
self.context_overlap = context_overlap
self.facemusk_dilation_ratio = facemusk_dilation_ratio
self.facecrop_dilation_ratio = facecrop_dilation_ratio
self.cfg = cfg
self.steps = steps
# Загрузка конфигураций и проверка устройства
self.config = OmegaConf.load(self.config)
self.weight_dtype = torch.float16 if self.config.weight_dtype == "fp16" else torch.float32
if self.device.__contains__("cuda") and not torch.cuda.is_available():
self.device = "cpu"
inference_config_path = self.config.inference_config
self.infer_config = OmegaConf.load(inference_config_path)
# Инициализация моделей
self._initialize_models()
# Установка параметров ширины и высоты
self.width, self.height = self.W, self.H
sched_kwargs = OmegaConf.to_container(self.infer_config.noise_scheduler_kwargs)
scheduler = DDIMScheduler(**sched_kwargs)
# Инициализация конвейера обработки видео
self.pipe = Audio2VideoPipeline(
vae=self.vae,
reference_unet=self.reference_unet,
denoising_unet=self.denoising_unet,
audio_guider=self.audio_processor,
face_locator=self.face_locator,
scheduler=scheduler,
).to("cuda", dtype=self.weight_dtype)
def _initialize_models(self):
"""
Инициализация моделей для обработки видео и аудио.
"""
# VAE модель
self.vae = AutoencoderKL.from_pretrained(self.config.pretrained_vae_path).to("cuda", dtype=self.weight_dtype)
# Reference UNet
self.reference_unet = UNet2DConditionModel.from_pretrained(self.config.pretrained_base_model_path, subfolder="unet").to(dtype=self.weight_dtype, device=self.device)
self.reference_unet.load_state_dict(torch.load(self.config.reference_unet_path, map_location="cpu"))
# Denoising UNet
self.denoising_unet = EchoUNet3DConditionModel.from_pretrained_2d(
self.config.pretrained_base_model_path,
self.config.motion_module_path if os.path.exists(self.config.motion_module_path) else "",
subfolder="unet",
unet_additional_kwargs=self.infer_config.unet_additional_kwargs
).to(dtype=self.weight_dtype, device=self.device)
self.denoising_unet.load_state_dict(torch.load(self.config.denoising_unet_path, map_location="cpu"), strict=False)
# Локатор лица
self.face_locator = FaceLocator(320, conditioning_channels=1, block_out_channels=(16, 32, 96, 256)).to(dtype=self.weight_dtype, device="cuda")
self.face_locator.load_state_dict(torch.load(self.config.face_locator_path))
# Аудиомодель
self.audio_processor = load_audio_model(model_path=self.config.audio_model_path, device=self.device)
# Детектор лица
self.face_detector = MTCNN(image_size=320, margin=0, min_face_size=20, thresholds=[0.6, 0.7, 0.7], factor=0.709, post_process=True, device=self.device)
def __call__(self, ref_image_path, audio_path, save_dir):
"""
Генерация видео на основе изображения и аудиофайла.
:param ref_image_path: Путь до изображения.
:param audio_path: Путь до аудиофайла.
"""
# Установка генератора случайных чисел для воспроизводимости
if self.seed is not None and self.seed > -1:
generator = torch.manual_seed(self.seed)
else:
generator = torch.manual_seed(random.randint(100, 1000000))
# Создание папки для сохранения выходных данных
os.makedirs(save_dir, exist_ok=True)
ref_name = Path(ref_image_path).stem
audio_name = Path(audio_path).stem
final_fps = self.fps
# Подготовка маски лица
face_img = cv2.imread(ref_image_path) # Загрузка референсного изображения
face_mask = np.zeros((face_img.shape[0], face_img.shape[1])).astype('uint8') # Инициализация маски лица
# Детектирование лиц на изображении
det_bboxes, probs = self.face_detector.detect(face_img)
select_bbox = self.select_face(det_bboxes, probs) # Выбор самого крупного лица
if select_bbox is None:
face_mask[:, :] = 255 # Если лицо не найдено, маска покрывает всё изображение
else:
# Вычисление координат и увеличение рамки лица
xyxy = select_bbox[:4]
xyxy = np.round(xyxy).astype('int')
rb, re, cb, ce = xyxy[1], xyxy[3], xyxy[0], xyxy[2]
r_pad = int((re - rb) * self.facemusk_dilation_ratio)
c_pad = int((ce - cb) * self.facemusk_dilation_ratio)
face_mask[rb - r_pad:re + r_pad, cb - c_pad:ce + c_pad] = 255
# Обрезка области лица
r_pad_crop = int((re - rb) * self.facecrop_dilation_ratio)
c_pad_crop = int((ce - cb) * self.facecrop_dilation_ratio)
crop_rect = [
max(0, cb - c_pad_crop),
max(0, rb - r_pad_crop),
min(ce + c_pad_crop, face_img.shape[1]),
min(re + r_pad_crop, face_img.shape[0])
]
face_img, _ = crop_and_pad(face_img, crop_rect) # Обрезка и дополнение до размеров
face_mask, _ = crop_and_pad(face_mask, crop_rect)
face_img = cv2.resize(face_img, (self.W, self.H)) # Изменение размера изображения лица
face_mask = cv2.resize(face_mask, (self.W, self.H)) # Изменение размера маски лица
# Конвертация лица в формат PIL
ref_image_pil = Image.fromarray(face_img[:, :, [2, 1, 0]]) # Перевод в RGB
# Преобразование маски в тензор
face_mask_tensor = torch.Tensor(face_mask).to(dtype=self.weight_dtype, device="cuda").unsqueeze(0).unsqueeze(0).unsqueeze(0) / 255.0
# Генерация видео с использованием конвейера
video = self.pipe(
ref_image_pil,
audio_path,
face_mask_tensor,
self.width,
self.height,
self.L,
self.steps,
self.cfg,
generator=generator,
audio_sample_rate=self.sample_rate,
context_frames=self.context_frames,
fps=final_fps,
context_overlap=self.context_overlap
).videos
# Сохранение видео в файл
save_videos_grid(
video,
f"{save_dir}/wish.mp4",
n_rows=1,
fps=final_fps,
)
# Добавление аудио к видео
video_clip = VideoFileClip(f"{save_dir}/wish.mp4")
audio_clip = AudioFileClip(audio_path)
video_clip = video_clip.set_audio(audio_clip)
video_clip.write_videofile(f"{save_dir}/wish_withaudio.mp4", codec="libx264", audio_codec="aac")
print(f"{save_dir}/wish_withaudio.mp4")
def select_face(self, det_bboxes, probs):
"""
Выбор самого крупного лица из списка детектированных лиц с вероятностью выше 0.8.
:param det_bboxes: Список координат ограничивающих рамок лиц.
:param probs: Список вероятностей для каждой рамки.
:return: Координаты самой крупной рамки или None, если лицо не найдено.
"""
if det_bboxes is None or probs is None:
return None
filtered_bboxes = [bbox for bbox, prob in zip(det_bboxes, probs) if prob > 0.8]
if not filtered_bboxes:
return None
# Сортировка по площади рамки (в порядке убывания)
sorted_bboxes = sorted(filtered_bboxes, key=lambda x: (x[3] - x[1]) * (x[2] - x[0]), reverse=True)
return sorted_bboxes[0]Вызов модели
from audio2video import Audio2Video
# Инициализация модели Audio2Video
a2v = Audio2Video()
# Создание видео из аудио
a2v(os.path.join("avatars", f"{character}.jpeg"), "wish.wav", "wishes")Пример данных:
Вот и всё! Теперь у нас есть готовая новогодняя ИИ‑открытка!
