Add ComfyUI HighSync nodes
This commit is contained in:
magicwang1111
@@ -0,0 +1,8 @@
|
||||
* text=auto eol=lf
|
||||
|
||||
*.png binary
|
||||
*.jpg binary
|
||||
*.jpeg binary
|
||||
*.npz binary
|
||||
*.tflite binary
|
||||
*.task binary
|
||||
@@ -0,0 +1,7 @@
|
||||
__pycache__/
|
||||
*.py[cod]
|
||||
.pytest_cache/
|
||||
.venv/
|
||||
cache/
|
||||
pretrained_weights/
|
||||
outputs/
|
||||
@@ -0,0 +1,47 @@
|
||||
# ComfyUI HighSync Wrapper
|
||||
|
||||
HighSync lip-sync nodes for ComfyUI. This wrapper uses ComfyUI `IMAGE` frames plus native `AUDIO`, runs the HighSync pipeline in-process, and returns processed frames, audio, output fps, and an encoded mp4 path.
|
||||
|
||||
## Nodes
|
||||
|
||||
- `HighSync Download Models`: checks or downloads the required model files into `D:/ComfyUI/models/highsync`.
|
||||
- `HighSync Model Loader`: loads and caches the VAE, UNets, scheduler, and Whisper audio processor.
|
||||
- `HighSync LipSync`: simple node for `IMAGE + AUDIO` input.
|
||||
- `HighSync LipSync Advanced`: exposes inference steps, CFG scale, seed, chunk size, denoising, trimming, and intermediate-output options.
|
||||
|
||||
## Model Layout
|
||||
|
||||
The default model directory is:
|
||||
|
||||
```text
|
||||
D:/ComfyUI/models/highsync
|
||||
```
|
||||
|
||||
It must contain:
|
||||
|
||||
```text
|
||||
denoising_unet-500.pth
|
||||
reference_unet-500.pth
|
||||
sd-vae-ft-mse/
|
||||
sd-image-variations-diffusers/
|
||||
audio_processor/whisper_tiny.pt
|
||||
```
|
||||
|
||||
Use `HighSync Download Models` with `download=True`, or download `saeed-5959/high_sync` from Hugging Face into that folder.
|
||||
|
||||
## Usage
|
||||
|
||||
1. Install `requirements.txt` into the same Python environment that runs ComfyUI.
|
||||
2. Restart ComfyUI.
|
||||
3. Use VideoHelperSuite or another loader to provide video frames as `IMAGE` and audio as `AUDIO`.
|
||||
4. Connect `HighSync Model Loader` to `HighSync LipSync`.
|
||||
|
||||
The wrapper converts video frames to 25fps and audio to mono 16kHz for HighSync inference. The returned frame rate is always `25` in this v1 implementation.
|
||||
|
||||
## Notes
|
||||
|
||||
- `max_frames=0` means process all available frames after 25fps conversion.
|
||||
- `denoise_audio=False` is the default because the official `denoiser` package pins older Hydra/OmegaConf versions than current ComfyUI environments. To enable it, install the package without pulling its old dependencies: `pip install --no-deps denoiser==0.1.5`.
|
||||
- When `denoise_audio=True`, the official denoiser is loaded lazily and cached. The final mp4 still uses the adjusted original audio track to preserve voice quality.
|
||||
- DeepFace/TensorFlow are optional. The runtime tries DeepFace if it is already installed, otherwise it falls back to OpenCV Haar face detection.
|
||||
- Outputs are written under `ComfyUI/output/highsync/<task_id>/final.mp4`.
|
||||
@@ -0,0 +1,6 @@
|
||||
try:
|
||||
from .nodes import NODE_CLASS_MAPPINGS, NODE_DISPLAY_NAME_MAPPINGS
|
||||
except ImportError:
|
||||
from nodes import NODE_CLASS_MAPPINGS, NODE_DISPLAY_NAME_MAPPINGS
|
||||
|
||||
__all__ = ["NODE_CLASS_MAPPINGS", "NODE_DISPLAY_NAME_MAPPINGS"]
|
||||
@@ -0,0 +1,11 @@
|
||||
.venv/
|
||||
caache/
|
||||
data/
|
||||
dataset/
|
||||
exp_output/
|
||||
out_ref_500/
|
||||
pretrained_weights/
|
||||
runs/
|
||||
test/
|
||||
__pycache__/
|
||||
cache/
|
||||
@@ -0,0 +1,138 @@
|
||||
<h1 align='center'>HighSync: High-Quality Lip Synchronization via
|
||||
Latent Diffusion Models</h1>
|
||||
|
||||
<div align='center'>
|
||||
<a href='https://github.com/saeed5959' target='_blank'>Saeed Firouzi</a><sup>1</sup> 
|
||||
</div>
|
||||
|
||||
<br>
|
||||
|
||||
<div align='center'>
|
||||
<a href='https://huggingface.co/saeed-5959/high_sync'><img src='https://img.shields.io/badge/%F0%9F%A4%97%20HuggingFace-Model-yellow'></a>
|
||||
<a href='https://arxiv.org/abs/2605.16918'><img src='https://img.shields.io/badge/Paper-Arxiv-red'></a>
|
||||
<a href='https://huggingface.co/datasets/saeed-5959/vfhq'><img src='https://img.shields.io/badge/Dataset-Hugging_Face-CFAFD4'></a>
|
||||
</div>
|
||||
|
||||
|
||||
## Abstraction
|
||||
We present HighSync, an end-to-end diffusion-based
|
||||
framework for high-fidelity lip synchronization that generates
|
||||
photorealistic talking-face videos aligned with arbitrary input
|
||||
audio. Existing approaches consistently struggle to reconcile
|
||||
image quality with synchronization accuracy, producing either
|
||||
visually degraded outputs or temporally inconsistent lip move-
|
||||
ments. HighSync addresses both challenges simultaneously and,
|
||||
to our knowledge, is the first lip sync model to operate natively
|
||||
at 512×512 resolution, positioning it as a viable solution for
|
||||
professional production environments such as the film and broad-
|
||||
cast industries. Central to our approach is the identification and
|
||||
systematic elimination of a data leakage phenomenon that has
|
||||
silently undermined temporal modeling in prior work, preventing
|
||||
models from developing a genuine dependence on the audio
|
||||
signal. Comprehensive evaluations across both perceptual quality
|
||||
and synchronization accuracy metrics confirm that HighSync
|
||||
achieves state-of-the-art performance on both fronts.
|
||||
|
||||
|
||||
## Model Structure
|
||||
<img src='./imgs/model_stage2.png'>
|
||||
|
||||
|
||||
## ⚒️ Installation
|
||||
|
||||
### Environment
|
||||
|
||||
Ubuntu 20 or 22
|
||||
|
||||
### Download the Codes
|
||||
|
||||
```bash
|
||||
git clone https://github.com/saeed5959/high_sync
|
||||
cd high_sync
|
||||
```
|
||||
|
||||
|
||||
### Install packages with `pip`
|
||||
```bash
|
||||
pip install -r requirements.txt
|
||||
```
|
||||
|
||||
### Install ffmpeg
|
||||
```bash
|
||||
apt-get install ffmpeg
|
||||
```
|
||||
|
||||
### Download pretrained weights
|
||||
|
||||
```shell
|
||||
git lfs install
|
||||
git clone https://huggingface.co/saeed-5959/high_sync pretrained_weights
|
||||
```
|
||||
|
||||
The **pretrained_weights** is organized as follows.
|
||||
|
||||
```
|
||||
./pretrained_weights/
|
||||
├── denoising_unet-500.pth
|
||||
├── reference_unet-500.pth
|
||||
├── sd-vae-ft-mse
|
||||
│ └── ...
|
||||
├── sd-image-variations-diffusers
|
||||
│ └── ...
|
||||
└── audio_processor
|
||||
└── whisper_tiny.pt
|
||||
```
|
||||
|
||||
In which **denoising_unet.pth** / **reference_unet.pth** / are the main checkpoints of **Highsync**. Other models in this hub can be also downloaded from it's original hub, thanks to their brilliant works:
|
||||
- [sd-vae-ft-mse](https://huggingface.co/stabilityai/sd-vae-ft-mse)
|
||||
- [sd-image-variations-diffusers](https://huggingface.co/lambdalabs/sd-image-variations-diffusers)
|
||||
- [audio_processor(whisper)](https://openaipublic.azureedge.net/main/whisper/models/65147644a518d12f04e32d6f3b26facc3f8dd46e5390956a9424a650c0ce22b9/tiny.pt)
|
||||
|
||||
### Inference
|
||||
|
||||
1)First convert your video to fps=25
|
||||
|
||||
```bash
|
||||
ffmpeg -i input.mp4 -r 25 out_25.mp4
|
||||
```
|
||||
|
||||
2)Then run the python inference script:
|
||||
|
||||
```bash
|
||||
python -m inference --source_video "video_path.mp4" --driving_audio "audio_path.wav" --output "save_path.mp4"
|
||||
```
|
||||
|
||||
### Dataset
|
||||
We preprocessed 3 public datasets and put their clean videos in these links :
|
||||
|
||||
[VFHQ](https://huggingface.co/datasets/saeed-5959/vfhq)
|
||||
|
||||
[Celebv-HQ](https://huggingface.co/datasets/saeed-5959/celebv_hq_head_talking)
|
||||
|
||||
[HDTF](https://huggingface.co/datasets/saeed-5959/hdtf)
|
||||
|
||||
Notice : thses videos has been preprocessed based on the paper appraoch!
|
||||
|
||||
## 🙏🏻 Acknowledgements
|
||||
|
||||
This work is mainly based on [EchoMimic](https://github.com/antgroup/echomimic) work.
|
||||
|
||||
We would like to thank the contributors to the [EchoMimic](), [AnimateDiff](https://github.com/guoyww/AnimateDiff), [Moore-AnimateAnyone](https://github.com/MooreThreads/Moore-AnimateAnyone) and [MuseTalk](https://github.com/TMElyralab/MuseTalk) repositories, for their open research and exploration.
|
||||
|
||||
We are also grateful to [V-Express](https://github.com/tencent-ailab/V-Express) and [hallo](https://github.com/fudan-generative-vision/hallo) for their outstanding work in the area of diffusion-based talking heads.
|
||||
|
||||
If we missed any open-source projects or related articles, we would like to complement the acknowledgement of this specific work immediately.
|
||||
|
||||
## Citation
|
||||
```bibtex
|
||||
@article{daghigh2024highsync,
|
||||
title={HighSync: High-Quality Lip Synchronization via Latent Diffusion Models},
|
||||
author={Saeed Firouzi Daghigh and Majid Iranpour Mobarekeh and Mostafa Alavi and Mehdi Bagheri},
|
||||
journal={arXiv preprint arXiv:2605.16918},
|
||||
year={2026}
|
||||
}
|
||||
```
|
||||
|
||||
|
||||
## 🌟 Star History
|
||||
[](https://star-history.com/#saeed_5959/highsync&Date)
|
||||
@@ -0,0 +1 @@
|
||||
"""HighSync research code packaged for the ComfyUI wrapper."""
|
||||
@@ -0,0 +1,72 @@
|
||||
mask_path: "./masks/mask_64.png"
|
||||
mask_512_path: "./masks/mask_512.png"
|
||||
|
||||
weight_dtype: fp16
|
||||
|
||||
data:
|
||||
n_motion_frames: 2
|
||||
n_sample_frames: 12
|
||||
source_image:
|
||||
width: 512
|
||||
height: 512
|
||||
driving_audio:
|
||||
sample_rate: 16000
|
||||
export_video:
|
||||
fps: 25
|
||||
|
||||
inference_steps: 20
|
||||
cfg_scale: 3.5
|
||||
|
||||
denoising_unet_path: "./pretrained_weights/denoising_unet-500.pth"
|
||||
reference_unet_path: "./pretrained_weights/reference_unet-500.pth"
|
||||
|
||||
base_model_path: "./pretrained_weights/sd-image-variations-diffusers/"
|
||||
|
||||
motion_module_path: ""
|
||||
|
||||
audio_model_path: "./pretrained_weights/audio_processor/whisper_tiny.pt"
|
||||
|
||||
vae:
|
||||
model_path: "./pretrained_weights/sd-vae-ft-mse/"
|
||||
|
||||
save_path: ./cache
|
||||
|
||||
unet_additional_kwargs:
|
||||
use_inflated_groupnorm: true
|
||||
unet_use_cross_frame_attention: false
|
||||
unet_use_temporal_attention: false
|
||||
use_motion_module: true
|
||||
cross_attention_dim: 384
|
||||
motion_module_resolutions:
|
||||
- 1
|
||||
- 2
|
||||
- 4
|
||||
- 8
|
||||
motion_module_mid_block: true
|
||||
motion_module_decoder_only: false
|
||||
motion_module_type: Vanilla
|
||||
motion_module_kwargs:
|
||||
num_attention_heads: 8
|
||||
num_transformer_block: 1
|
||||
attention_block_types:
|
||||
- Temporal_Self
|
||||
- Temporal_Self
|
||||
temporal_position_encoding: true
|
||||
temporal_position_encoding_max_len: 32
|
||||
temporal_attention_dim_div: 1
|
||||
|
||||
|
||||
enable_zero_snr: true
|
||||
|
||||
noise_scheduler_kwargs:
|
||||
beta_start: 0.00085
|
||||
beta_end: 0.012
|
||||
beta_schedule: "linear"
|
||||
clip_sample: false
|
||||
steps_offset: 1
|
||||
### Zero-SNR params
|
||||
prediction_type: "v_prediction"
|
||||
rescale_betas_zero_snr: True
|
||||
timestep_spacing: "trailing"
|
||||
|
||||
sampler: DDIM
|
||||
@@ -0,0 +1,60 @@
|
||||
data:
|
||||
train_bs: 16
|
||||
train_width: 512
|
||||
train_height: 512
|
||||
train_path: "./data/echo_train.json"
|
||||
val_path: "./data/echo_train.json"
|
||||
# Margin of frame indexes between ref and tgt images
|
||||
sample_margin: 30
|
||||
mask_path: "./masks/mask_64.png"
|
||||
mask_512_path: "./masks/mask_512.png"
|
||||
|
||||
solver:
|
||||
gradient_accumulation_steps: 1
|
||||
mixed_precision: "no"
|
||||
enable_xformers_memory_efficient_attention: True
|
||||
gradient_checkpointing: False
|
||||
max_train_steps: 300000
|
||||
max_grad_norm: 1.0
|
||||
# lr
|
||||
learning_rate: 1.0e-5
|
||||
scale_lr: False
|
||||
lr_warmup_steps: 100
|
||||
lr_scheduler: "constant"
|
||||
|
||||
# optimizer
|
||||
use_8bit_adam: False
|
||||
adam_beta1: 0.9
|
||||
adam_beta2: 0.999
|
||||
adam_weight_decay: 1.0e-2
|
||||
adam_epsilon: 1.0e-8
|
||||
|
||||
val:
|
||||
validation_steps: 1000
|
||||
|
||||
noise_scheduler_kwargs:
|
||||
num_train_timesteps: 1000
|
||||
beta_start: 0.00085
|
||||
beta_end: 0.012
|
||||
beta_schedule: "scaled_linear"
|
||||
steps_offset: 1
|
||||
clip_sample: false
|
||||
|
||||
base_model_path: "./pretrained_weights/sd-image-variations-diffusers/"
|
||||
vae_model_path: "./pretrained_weights/sd-vae-ft-mse"
|
||||
face_analysis_model_path: "./pretrained_models/face_analysis"
|
||||
echo_path: "./exp_output/stage1"
|
||||
|
||||
weight_dtype: "fp16" # [fp16, fp32]
|
||||
uncond_ratio: 0.1
|
||||
noise_offset: 0.05
|
||||
snr_gamma: 5.0
|
||||
enable_zero_snr: True
|
||||
face_locator_pretrained: False
|
||||
cross_attention_dim: 384
|
||||
|
||||
seed: 42
|
||||
resume_from_checkpoint: "latest"
|
||||
checkpointing_steps: 10000
|
||||
exp_name: "stage1"
|
||||
output_dir: "./exp_output"
|
||||
@@ -0,0 +1,109 @@
|
||||
data:
|
||||
train_bs: 4
|
||||
val_bs: 1
|
||||
train_width: 512
|
||||
train_height: 512
|
||||
fps: 25
|
||||
sample_rate: 16000
|
||||
n_sample_frames: 12
|
||||
audio_margin: 2
|
||||
train_path: "./data/echo_train.json"
|
||||
val_path: "./data/echo_train.json"
|
||||
mask_path: "./masks/mask_64.png"
|
||||
mask_512_path: "./masks/mask_512.png"
|
||||
|
||||
wav2vec_config:
|
||||
audio_type: "vocals" # audio vocals
|
||||
model_scale: "base" # base large
|
||||
features: "all" # last avg all
|
||||
model_path: ./pretrained_models/wav2vec/wav2vec2-base-960h
|
||||
audio_separator:
|
||||
model_path: ./pretrained_models/audio_separator/Kim_Vocal_2.onnx
|
||||
face_expand_ratio: 1.2
|
||||
|
||||
solver:
|
||||
gradient_accumulation_steps: 1
|
||||
mixed_precision: "no"
|
||||
enable_xformers_memory_efficient_attention: True
|
||||
gradient_checkpointing: True
|
||||
max_train_steps: 100000
|
||||
max_grad_norm: 1.0
|
||||
# lr
|
||||
learning_rate: 5e-6
|
||||
scale_lr: False
|
||||
lr_warmup_steps: 1
|
||||
lr_scheduler: "constant"
|
||||
|
||||
# optimizer
|
||||
use_8bit_adam: True
|
||||
adam_beta1: 0.9
|
||||
adam_beta2: 0.999
|
||||
adam_weight_decay: 1.0e-2
|
||||
adam_epsilon: 1.0e-8
|
||||
|
||||
val:
|
||||
validation_steps: 1000
|
||||
|
||||
noise_scheduler_kwargs:
|
||||
num_train_timesteps: 1000
|
||||
beta_start: 0.00085
|
||||
beta_end: 0.012
|
||||
beta_schedule: "linear"
|
||||
steps_offset: 1
|
||||
clip_sample: false
|
||||
|
||||
unet_additional_kwargs:
|
||||
use_inflated_groupnorm: true
|
||||
unet_use_cross_frame_attention: false
|
||||
unet_use_temporal_attention: false
|
||||
use_motion_module: true
|
||||
cross_attention_dim: 384
|
||||
motion_module_resolutions:
|
||||
- 1
|
||||
- 2
|
||||
- 4
|
||||
- 8
|
||||
motion_module_mid_block: true
|
||||
motion_module_decoder_only: false
|
||||
motion_module_type: Vanilla
|
||||
motion_module_kwargs:
|
||||
num_attention_heads: 8
|
||||
num_transformer_block: 1
|
||||
attention_block_types:
|
||||
- Temporal_Self
|
||||
- Temporal_Self
|
||||
temporal_position_encoding: true
|
||||
temporal_position_encoding_max_len: 32
|
||||
temporal_attention_dim_div: 1
|
||||
|
||||
trainable_para:
|
||||
- motion_modules
|
||||
|
||||
base_model_path: "./pretrained_weights/sd-image-variations-diffusers/"
|
||||
vae_model_path: "./pretrained_weights/sd-vae-ft-mse"
|
||||
face_analysis_model_path: "./pretrained_models/face_analysis"
|
||||
mm_path: "./pretrained_weights/motion_module.pth"
|
||||
echo_path: "./exp_output" #"./pretrained_weights"
|
||||
resume_denoise_path: "./exp_output/denoising_unet-44000.pth"
|
||||
resume_reference_path: "./exp_output/reference_unet-44000.pth"
|
||||
resume: true
|
||||
|
||||
weight_dtype: "fp16" # [fp16, fp32]
|
||||
uncond_img_ratio: 0.1
|
||||
uncond_audio_ratio: 0.0
|
||||
uncond_target_ratio: 0.0
|
||||
start_ratio: 1.0
|
||||
noise_offset: 0.05
|
||||
snr_gamma: 5.0
|
||||
enable_zero_snr: True
|
||||
stage1_ckpt_dir: "./exp_output"
|
||||
|
||||
single_inference_times: 10
|
||||
inference_steps: 20
|
||||
cfg_scale: 3.5
|
||||
|
||||
seed: 42
|
||||
resume_from_checkpoint: "latest"
|
||||
checkpointing_steps: 2000
|
||||
exp_name: "stage2"
|
||||
output_dir: "./exp_output"
|
||||
@@ -0,0 +1,157 @@
|
||||
# pylint: disable=W1203,W0718
|
||||
"""
|
||||
This module is used to process videos to prepare data for training. It utilizes various libraries and models
|
||||
to perform tasks such as video frame extraction, audio extraction, face mask generation, and face embedding extraction.
|
||||
The script takes in command-line arguments to specify the input and output directories, GPU status, level of parallelism,
|
||||
and rank for distributed processing.
|
||||
|
||||
Usage:
|
||||
python -m scripts.data_preprocess --input_dir /path/to/video_dir --dataset_name dataset_name --gpu_status --parallelism 4 --rank 0
|
||||
|
||||
Example:
|
||||
python -m scripts.data_preprocess -i data/videos -o data/output -g -p 4 -r 0
|
||||
"""
|
||||
import argparse
|
||||
import logging
|
||||
import os
|
||||
from pathlib import Path
|
||||
from typing import List
|
||||
import soundfile as sf
|
||||
import torchaudio
|
||||
from denoiser import pretrained
|
||||
from denoiser.dsp import convert_audio
|
||||
import torch
|
||||
from tqdm import tqdm
|
||||
|
||||
from .src.models.whisper.audio2feature import Audio2Feature
|
||||
from .src.datasets.image_processor import ImageProcessorForDataProcessing
|
||||
from .src.utils.util import convert_video_to_images, extract_audio_from_videos
|
||||
|
||||
# Configure logging
|
||||
logging.basicConfig(level=logging.INFO,
|
||||
format='%(asctime)s - %(levelname)s - %(message)s')
|
||||
|
||||
|
||||
# model = pretrained.dns64().cuda()
|
||||
|
||||
def process_audio(input_audio_path):
|
||||
|
||||
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(input_audio_path, denoised.data.cpu().numpy()[0], 16000)
|
||||
|
||||
return
|
||||
|
||||
|
||||
def process_single_video(video_path: Path,
|
||||
output_dir: Path,
|
||||
audio_processor: Audio2Feature,
|
||||
step: int) -> None:
|
||||
"""
|
||||
Process a single video file.
|
||||
|
||||
Args:
|
||||
video_path (Path): Path to the video file.
|
||||
output_dir (Path): Directory to save the output.
|
||||
audio_processor (AudioProcessor): Audio processor object.
|
||||
gpu_status (bool): Whether to use GPU for processing.
|
||||
"""
|
||||
assert video_path.exists(), f"Video path {video_path} does not exist"
|
||||
|
||||
audio_emb_dir = output_dir / "audio_emb_whisper"
|
||||
audio_emb_dir.mkdir(parents=True, exist_ok=True)
|
||||
|
||||
# try:
|
||||
if step == 1:
|
||||
images_output_dir = output_dir / 'images' / video_path.stem
|
||||
images_output_dir.mkdir(parents=True, exist_ok=True)
|
||||
images_output_dir = convert_video_to_images(
|
||||
video_path, images_output_dir)
|
||||
logging.info(f"Images saved to: {images_output_dir}")
|
||||
|
||||
audio_output_dir = output_dir / 'audios'
|
||||
audio_output_dir.mkdir(parents=True, exist_ok=True)
|
||||
audio_output_path = audio_output_dir / f'{video_path.stem}.wav'
|
||||
audio_output_path = extract_audio_from_videos(
|
||||
video_path, audio_output_path)
|
||||
logging.info(f"Audio extracted to: {audio_output_path}")
|
||||
|
||||
else:
|
||||
audio_path = output_dir / "audios" / f"{video_path.stem}.wav"
|
||||
# process_audio(str(audio_path))
|
||||
audio_emb = audio_processor.audio2feat(str(audio_path))
|
||||
audio_emb = audio_processor.feature2chunks(feature_array=audio_emb, fps=25)
|
||||
torch.save(audio_emb, str(
|
||||
audio_emb_dir / f"{video_path.stem}.pt"))
|
||||
# except Exception as e:
|
||||
# logging.error(f"Failed to process video {video_path}: {e}")
|
||||
|
||||
|
||||
def process_all_videos(input_video_list: List[Path], output_dir: Path, step: int) -> None:
|
||||
"""
|
||||
Process all videos in the input list.
|
||||
|
||||
Args:
|
||||
input_video_list (List[Path]): List of video paths to process.
|
||||
output_dir (Path): Directory to save the output.
|
||||
gpu_status (bool): Whether to use GPU for processing.
|
||||
"""
|
||||
whisper_model_path = "./pretrained_weights/audio_processor/whisper_tiny.pt"
|
||||
|
||||
audio_processor = Audio2Feature(model_path=whisper_model_path, device="cuda") if step==2 else None
|
||||
|
||||
for video_path in tqdm(input_video_list, desc="Processing videos"):
|
||||
process_single_video(video_path, output_dir, audio_processor, step)
|
||||
|
||||
|
||||
def get_video_paths(source_dir: Path, parallelism: int, rank: int) -> List[Path]:
|
||||
"""
|
||||
Get paths of videos to process, partitioned for parallel processing.
|
||||
|
||||
Args:
|
||||
source_dir (Path): Source directory containing videos.
|
||||
parallelism (int): Level of parallelism.
|
||||
rank (int): Rank for distributed processing.
|
||||
|
||||
Returns:
|
||||
List[Path]: List of video paths to process.
|
||||
"""
|
||||
video_paths = [item for item in sorted(
|
||||
source_dir.iterdir()) if item.is_file() and item.suffix == '.mp4']
|
||||
return [video_paths[i] for i in range(len(video_paths)) if i % parallelism == rank]
|
||||
|
||||
|
||||
if __name__ == "__main__":
|
||||
parser = argparse.ArgumentParser(
|
||||
description="Process videos to prepare data for training. Run this script twice with different GPU status parameters."
|
||||
)
|
||||
parser.add_argument("-i", "--input_dir", type=Path,
|
||||
required=True, help="Directory containing videos")
|
||||
parser.add_argument("-o", "--output_dir", type=Path,
|
||||
help="Directory to save results, default is parent dir of input dir")
|
||||
parser.add_argument("-s", "--step", type=int, default=1,
|
||||
help="Specify data processing step 1 or 2, you should run 1 and 2 sequently")
|
||||
parser.add_argument("-p", "--parallelism", default=1,
|
||||
type=int, help="Level of parallelism")
|
||||
parser.add_argument("-r", "--rank", default=0, type=int,
|
||||
help="Rank for distributed processing")
|
||||
|
||||
args = parser.parse_args()
|
||||
|
||||
if args.output_dir is None:
|
||||
args.output_dir = args.input_dir.parent
|
||||
|
||||
video_path_list = get_video_paths(
|
||||
args.input_dir, args.parallelism, args.rank)
|
||||
|
||||
videos_out_path = os.path.join(args.output_dir, "videos")
|
||||
# os.system(f"cp -r {args.input_dir} {videos_out_path}")
|
||||
|
||||
if not video_path_list:
|
||||
logging.warning("No videos to process.")
|
||||
else:
|
||||
process_all_videos(video_path_list, args.output_dir, args.step)
|
||||
@@ -0,0 +1,460 @@
|
||||
# pylint: disable=W0718
|
||||
"""
|
||||
This module is responsible for processing images, particularly for face-related tasks.
|
||||
It uses various libraries such as OpenCV, NumPy, and InsightFace to perform tasks like
|
||||
face detection, augmentation, and mask rendering. The ImageProcessor class encapsulates
|
||||
the functionality for these operations.
|
||||
"""
|
||||
import os
|
||||
from typing import List
|
||||
|
||||
import cv2
|
||||
import mediapipe as mp
|
||||
import numpy as np
|
||||
import torch
|
||||
from insightface.app import FaceAnalysis
|
||||
from PIL import Image, ImageOps
|
||||
from torchvision import transforms
|
||||
|
||||
|
||||
MEAN = 0.5
|
||||
STD = 0.5
|
||||
|
||||
|
||||
def _select_largest_deepface_bbox(faces):
|
||||
max_face = 0
|
||||
best_bbox = (0, 0, 0, 0)
|
||||
for face in faces:
|
||||
facial_area = face.get("facial_area") or {}
|
||||
face_area = facial_area.get("w", 0) * facial_area.get("h", 0)
|
||||
if face.get("confidence", 0) > 0.5 and face_area > max_face:
|
||||
best_bbox = (
|
||||
int(facial_area["x"]),
|
||||
int(facial_area["y"]),
|
||||
int(facial_area["x"] + facial_area["w"]),
|
||||
int(facial_area["y"] + facial_area["h"]),
|
||||
)
|
||||
max_face = face_area
|
||||
return best_bbox
|
||||
|
||||
|
||||
def _detect_face_with_deepface(frame):
|
||||
try:
|
||||
from deepface import DeepFace
|
||||
except Exception:
|
||||
return None
|
||||
|
||||
try:
|
||||
faces = DeepFace.extract_faces(
|
||||
img_path=frame,
|
||||
detector_backend="yolov8",
|
||||
align=False,
|
||||
enforce_detection=False,
|
||||
expand_percentage=0,
|
||||
)
|
||||
except Exception as exc:
|
||||
print(f"DeepFace detection failed, falling back to OpenCV Haar detection: {exc}")
|
||||
return None
|
||||
|
||||
if not faces:
|
||||
return (0, 0, 0, 0)
|
||||
return _select_largest_deepface_bbox(faces)
|
||||
|
||||
|
||||
def _detect_face_with_opencv(frame):
|
||||
cascade_path = cv2.data.haarcascades + "haarcascade_frontalface_default.xml"
|
||||
detector = cv2.CascadeClassifier(cascade_path)
|
||||
if detector.empty():
|
||||
return (0, 0, 0, 0)
|
||||
|
||||
gray = cv2.cvtColor(frame, cv2.COLOR_BGR2GRAY)
|
||||
faces = detector.detectMultiScale(gray, scaleFactor=1.1, minNeighbors=5, minSize=(40, 40))
|
||||
if len(faces) == 0:
|
||||
return (0, 0, 0, 0)
|
||||
|
||||
x, y, w, h = max(faces, key=lambda item: item[2] * item[3])
|
||||
return int(x), int(y), int(x + w), int(y + h)
|
||||
|
||||
|
||||
def detect_face_bbox(frame):
|
||||
bbox = _detect_face_with_deepface(frame)
|
||||
if bbox is None:
|
||||
bbox = _detect_face_with_opencv(frame)
|
||||
if bbox == (0, 0, 0, 0):
|
||||
print("No faces detected in the image. Using the entire image as the face region.")
|
||||
return bbox
|
||||
|
||||
class ImageProcessor:
|
||||
"""
|
||||
ImageProcessor is a class responsible for processing images, particularly for face-related tasks.
|
||||
It takes in an image and performs various operations such as augmentation, face detection,
|
||||
face embedding extraction, and rendering a face mask. The processed images are then used for
|
||||
further analysis or recognition purposes.
|
||||
|
||||
Attributes:
|
||||
img_size (int): The size of the image to be processed.
|
||||
|
||||
Methods:
|
||||
preprocess(source_image_path, cache_dir):
|
||||
Preprocesses the input image by performing augmentation, face detection,
|
||||
face embedding extraction, and rendering a face mask.
|
||||
|
||||
close():
|
||||
Closes the ImageProcessor and releases any resources being used.
|
||||
|
||||
_augmentation(images, transform, state=None):
|
||||
Applies image augmentation to the input images using the given transform and state.
|
||||
|
||||
__enter__():
|
||||
Enters a runtime context and returns the ImageProcessor object.
|
||||
|
||||
__exit__(_exc_type, _exc_val, _exc_tb):
|
||||
Exits a runtime context and handles any exceptions that occurred during the processing.
|
||||
"""
|
||||
def __init__(self, img_size) -> None:
|
||||
self.img_size = img_size
|
||||
|
||||
self.pixel_transform = transforms.Compose(
|
||||
[
|
||||
transforms.Resize(self.img_size),
|
||||
transforms.ToTensor(),
|
||||
transforms.Normalize([MEAN], [STD]),
|
||||
]
|
||||
)
|
||||
|
||||
def preprocess_frames(self, frames):
|
||||
"""
|
||||
Apply preprocessing to prepare for face analysis.
|
||||
|
||||
Parameters:
|
||||
source_image_path (str): The path to the source image.
|
||||
cache_dir (str): The directory to cache intermediate results.
|
||||
|
||||
Returns:
|
||||
None
|
||||
"""
|
||||
state = torch.get_rng_state()
|
||||
def find_max_h(bbox_list):
|
||||
|
||||
h_max = 0
|
||||
for bbox in bbox_list:
|
||||
h = bbox[3] - bbox[1]
|
||||
if h>h_max:
|
||||
h_max = h
|
||||
|
||||
return h_max
|
||||
|
||||
pixel_values_ref_img_list, bbox_list, bbox_final_list, frame_crop_list = [], [], [], []
|
||||
|
||||
for frame in frames:
|
||||
|
||||
bbox_list.append(detect_face_bbox(frame))
|
||||
|
||||
h_max = find_max_h(bbox_list)
|
||||
|
||||
for frame, bbox in zip(frames, bbox_list):
|
||||
|
||||
if bbox==(0,0,0,0):
|
||||
bbox_final_list.append((0,0,0,0))
|
||||
frame_crop = frame
|
||||
frame_crop_list.append(cv2.resize(frame_crop, (512,512)))
|
||||
else:
|
||||
frame_h, frame_w, _ = frame.shape
|
||||
|
||||
w, h = bbox[2] - bbox[0], bbox[3] - bbox[1]
|
||||
|
||||
if h_max < 1.4*h:
|
||||
h = h_max
|
||||
|
||||
bbox_final = (int(max(bbox[0] - 0.1*w, 0)), int(bbox[1]),
|
||||
int(min(bbox[2] + 0.1*w, frame_w)), int(min(bbox[1] + 1.1*h, frame_h)))
|
||||
# bbox_final = (0,0,512,512)
|
||||
bbox_final_list.append(bbox_final)
|
||||
|
||||
frame_face = frame[bbox_final[1]:bbox_final[3], bbox_final[0]:bbox_final[2]]
|
||||
frame_crop = cv2.resize(frame_face, (512,512))
|
||||
frame_crop_list.append(frame_crop)
|
||||
|
||||
pixel_values_ref_img_list.append(self._augmentation(Image.fromarray(cv2.cvtColor(frame_crop, cv2.COLOR_BGR2RGB)).convert("RGB"), self.pixel_transform, state).unsqueeze(0))
|
||||
|
||||
return pixel_values_ref_img_list, bbox_final_list, frame_crop_list
|
||||
|
||||
|
||||
def preprocess_frames_jump(self, frames):
|
||||
"""
|
||||
Apply preprocessing to prepare for face analysis.
|
||||
|
||||
Parameters:
|
||||
source_image_path (str): The path to the source image.
|
||||
cache_dir (str): The directory to cache intermediate results.
|
||||
|
||||
Returns:
|
||||
None
|
||||
"""
|
||||
frame_h, frame_w, _ = frames[0].shape
|
||||
|
||||
state = torch.get_rng_state()
|
||||
def find_max_h(bbox_list):
|
||||
|
||||
h_max = 0
|
||||
h_min = 100000
|
||||
for bbox in bbox_list:
|
||||
h = bbox[3] - bbox[1]
|
||||
if h>h_max:
|
||||
h_max = h
|
||||
if h<h_min:
|
||||
h_min = h
|
||||
|
||||
return h_max, h_min
|
||||
|
||||
def differ(bbox_list):
|
||||
|
||||
jumps = []
|
||||
count_prev = 0
|
||||
for count, bbox in enumerate(bbox_list):
|
||||
bbox_prev = bbox_list[max(count-1,0)]
|
||||
diff_x = ( np.abs(bbox[0] - bbox_prev[0]) + np.abs(bbox[2] - bbox_prev[2]) ) / frame_w
|
||||
diff_y = ( np.abs(bbox[1] - bbox_prev[1]) + np.abs(bbox[3] - bbox_prev[3]) ) / frame_h
|
||||
|
||||
if ( diff_x + diff_y ) > 0.2:
|
||||
jumps.append((count_prev, count-1))
|
||||
count_prev = count
|
||||
|
||||
jumps.append((count_prev, count))
|
||||
|
||||
return jumps
|
||||
|
||||
bbox_list = []
|
||||
|
||||
for frame in frames:
|
||||
|
||||
bbox_list.append(detect_face_bbox(frame))
|
||||
|
||||
pixel_values_ref_img_all, bbox_final_all, frames_list_chunk = [], [], []
|
||||
|
||||
jumps = differ(bbox_list)
|
||||
for jump in jumps:
|
||||
print(jump)
|
||||
pixel_values_ref_img_list, bbox_final_list = [], []
|
||||
|
||||
bbox_chunk = bbox_list[jump[0]:jump[1]+1]
|
||||
frames_chunk = frames[jump[0]:jump[1]+1]
|
||||
|
||||
h_max, h_min = find_max_h(bbox_chunk)
|
||||
use_h = h_max < 1.4*h_min
|
||||
print(f"use h : {use_h}")
|
||||
|
||||
for frame, bbox in zip(frames_chunk, bbox_chunk):
|
||||
|
||||
if bbox==(0,0,0,0):
|
||||
bbox_final_list.append((0,0,0,0))
|
||||
frame_crop = frame
|
||||
|
||||
else:
|
||||
w, h = bbox[2] - bbox[0], bbox[3] - bbox[1]
|
||||
|
||||
if use_h:
|
||||
h = h_max
|
||||
|
||||
bbox_final = (int(max(bbox[0] - 0.1*w, 0)), int(bbox[1]),
|
||||
int(min(bbox[2] + 0.1*w, frame_w)), int(min(bbox[1] + 1.1*h, frame_h)))
|
||||
|
||||
bbox_final_list.append(bbox_final)
|
||||
|
||||
frame_face = frame[bbox_final[1]:bbox_final[3], bbox_final[0]:bbox_final[2]]
|
||||
frame_crop = cv2.resize(frame_face, (512,512))
|
||||
|
||||
pixel_values_ref_img_list.append(self._augmentation(Image.fromarray(cv2.cvtColor(frame_crop, cv2.COLOR_BGR2RGB)).convert("RGB"), self.pixel_transform, state).unsqueeze(0))
|
||||
|
||||
pixel_values_ref_img_all.append(pixel_values_ref_img_list)
|
||||
bbox_final_all.append(bbox_final_list)
|
||||
frames_list_chunk.append(frames_chunk)
|
||||
|
||||
return pixel_values_ref_img_all, bbox_final_all, frames_list_chunk
|
||||
|
||||
|
||||
def preprocess_frames_deepface_no_detect(self, frames):
|
||||
"""
|
||||
Apply preprocessing to the 14 source images to prepare for face analysis.
|
||||
|
||||
Parameters:
|
||||
source_image_path (str): The path to the source image.
|
||||
cache_dir (str): The directory to cache intermediate results.
|
||||
|
||||
Returns:
|
||||
None
|
||||
"""
|
||||
pixel_values_ref_img_list = []
|
||||
for frame in frames:
|
||||
pixel_values_ref_img_list.append(self._augmentation(Image.fromarray(frame).convert("RGB"), self.pixel_transform).unsqueeze(0))
|
||||
|
||||
return pixel_values_ref_img_list
|
||||
|
||||
|
||||
|
||||
def close(self):
|
||||
"""
|
||||
Closes the ImageProcessor and releases any resources held by the FaceAnalysis instance.
|
||||
|
||||
Args:
|
||||
self: The ImageProcessor instance.
|
||||
|
||||
Returns:
|
||||
None.
|
||||
"""
|
||||
for _, model in self.face_analysis.models.items():
|
||||
if hasattr(model, "Dispose"):
|
||||
model.Dispose()
|
||||
|
||||
def _augmentation(self, images, transform, state=None):
|
||||
if state is not None:
|
||||
torch.set_rng_state(state)
|
||||
if isinstance(images, List):
|
||||
transformed_images = [transform(img) for img in images]
|
||||
ret_tensor = torch.stack(transformed_images, dim=0) # (f, c, h, w)
|
||||
else:
|
||||
ret_tensor = transform(images) # (c, h, w)
|
||||
return ret_tensor
|
||||
|
||||
def __enter__(self):
|
||||
return self
|
||||
|
||||
def __exit__(self, _exc_type, _exc_val, _exc_tb):
|
||||
self.close()
|
||||
|
||||
|
||||
class ImageProcessorForDataProcessing():
|
||||
"""
|
||||
ImageProcessor is a class responsible for processing images, particularly for face-related tasks.
|
||||
It takes in an image and performs various operations such as augmentation, face detection,
|
||||
face embedding extraction, and rendering a face mask. The processed images are then used for
|
||||
further analysis or recognition purposes.
|
||||
|
||||
Attributes:
|
||||
img_size (int): The size of the image to be processed.
|
||||
face_analysis_model_path (str): The path to the face analysis model.
|
||||
|
||||
Methods:
|
||||
preprocess(source_image_path, cache_dir):
|
||||
Preprocesses the input image by performing augmentation, face detection,
|
||||
face embedding extraction, and rendering a face mask.
|
||||
|
||||
close():
|
||||
Closes the ImageProcessor and releases any resources being used.
|
||||
|
||||
_augmentation(images, transform, state=None):
|
||||
Applies image augmentation to the input images using the given transform and state.
|
||||
|
||||
__enter__():
|
||||
Enters a runtime context and returns the ImageProcessor object.
|
||||
|
||||
__exit__(_exc_type, _exc_val, _exc_tb):
|
||||
Exits a runtime context and handles any exceptions that occurred during the processing.
|
||||
"""
|
||||
def __init__(self, face_analysis_model_path, landmark_model_path, step) -> None:
|
||||
if step == 2:
|
||||
self.face_analysis = FaceAnalysis(
|
||||
name="",
|
||||
root=face_analysis_model_path,
|
||||
providers=["CUDAExecutionProvider", "CPUExecutionProvider"],
|
||||
)
|
||||
self.face_analysis.prepare(ctx_id=0, det_size=(640, 640))
|
||||
self.landmarker = None
|
||||
else:
|
||||
BaseOptions = mp.tasks.BaseOptions
|
||||
FaceLandmarker = mp.tasks.vision.FaceLandmarker
|
||||
FaceLandmarkerOptions = mp.tasks.vision.FaceLandmarkerOptions
|
||||
VisionRunningMode = mp.tasks.vision.RunningMode
|
||||
# Create a face landmarker instance with the video mode:
|
||||
options = FaceLandmarkerOptions(
|
||||
base_options=BaseOptions(model_asset_path=landmark_model_path),
|
||||
running_mode=VisionRunningMode.IMAGE,
|
||||
)
|
||||
self.landmarker = None #FaceLandmarker.create_from_options(options)
|
||||
self.face_analysis = None
|
||||
|
||||
def preprocess(self, source_image_path: str):
|
||||
"""
|
||||
Apply preprocessing to the source image to prepare for face analysis.
|
||||
|
||||
Parameters:
|
||||
source_image_path (str): The path to the source image.
|
||||
cache_dir (str): The directory to cache intermediate results.
|
||||
|
||||
Returns:
|
||||
None
|
||||
"""
|
||||
# 1. get face embdeding
|
||||
face_mask, face_emb, sep_pose_mask, sep_face_mask, sep_lip_mask = None, None, None, None, None
|
||||
if self.face_analysis:
|
||||
frame = sorted(os.listdir(source_image_path))[0]
|
||||
source_image = Image.open(
|
||||
os.path.join(source_image_path, frame))
|
||||
ref_image_pil = source_image.convert("RGB")
|
||||
|
||||
padding = 250
|
||||
ref_image_pil = ImageOps.expand(ref_image_pil, border=padding, fill=0)
|
||||
|
||||
# 2.1 detect face
|
||||
faces = self.face_analysis.get(cv2.cvtColor(
|
||||
np.array(ref_image_pil.copy()), cv2.COLOR_RGB2BGR))
|
||||
|
||||
# use max size face
|
||||
face = sorted(faces, key=lambda x: (
|
||||
x["bbox"][2] - x["bbox"][0]) * (x["bbox"][3] - x["bbox"][1]))[-1]
|
||||
# 2.2 face embedding
|
||||
face_emb = face["embedding"]
|
||||
if face_emb is None:
|
||||
print(f"error in image embedding : {source_image_path}")
|
||||
raise
|
||||
|
||||
if self.landmarker:
|
||||
# 3.1 get landmark
|
||||
landmarks, height, width = get_landmark_overframes(
|
||||
self.landmarker, source_image_path)
|
||||
assert len(landmarks) == len(os.listdir(source_image_path))
|
||||
|
||||
# 3 render face and lip mask
|
||||
face_mask = get_union_face_mask(landmarks, height, width)
|
||||
lip_mask = get_union_lip_mask(landmarks, height, width)
|
||||
|
||||
# 4 gaussian blur
|
||||
blur_face_mask = blur_mask(face_mask, (64, 64), (51, 51))
|
||||
blur_lip_mask = blur_mask(lip_mask, (64, 64), (31, 31))
|
||||
|
||||
# 5 seperate mask
|
||||
sep_face_mask = cv2.subtract(blur_face_mask, blur_lip_mask)
|
||||
sep_pose_mask = 255.0 - blur_face_mask
|
||||
sep_lip_mask = blur_lip_mask
|
||||
|
||||
return face_mask, face_emb, sep_pose_mask, sep_face_mask, sep_lip_mask
|
||||
|
||||
def close(self):
|
||||
"""
|
||||
Closes the ImageProcessor and releases any resources held by the FaceAnalysis instance.
|
||||
|
||||
Args:
|
||||
self: The ImageProcessor instance.
|
||||
|
||||
Returns:
|
||||
None.
|
||||
"""
|
||||
for _, model in self.face_analysis.models.items():
|
||||
if hasattr(model, "Dispose"):
|
||||
model.Dispose()
|
||||
|
||||
def _augmentation(self, images, transform, state=None):
|
||||
if state is not None:
|
||||
torch.set_rng_state(state)
|
||||
if isinstance(images, List):
|
||||
transformed_images = [transform(img) for img in images]
|
||||
ret_tensor = torch.stack(transformed_images, dim=0) # (f, c, h, w)
|
||||
else:
|
||||
ret_tensor = transform(images) # (c, h, w)
|
||||
return ret_tensor
|
||||
|
||||
def __enter__(self):
|
||||
return self
|
||||
|
||||
def __exit__(self, _exc_type, _exc_val, _exc_tb):
|
||||
self.close()
|
||||
Binary file not shown.
|
After Width: | Height: | Size: 403 KiB |
@@ -0,0 +1,398 @@
|
||||
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)
|
||||
Binary file not shown.
|
After Width: | Height: | Size: 136 B |
Binary file not shown.
|
After Width: | Height: | Size: 205 B |
Binary file not shown.
|
After Width: | Height: | Size: 3.3 KiB |
Binary file not shown.
|
After Width: | Height: | Size: 345 B |
Binary file not shown.
|
After Width: | Height: | Size: 98 B |
@@ -0,0 +1,42 @@
|
||||
accelerate==0.28.0
|
||||
audio-separator==0.17.2
|
||||
av==12.1.0
|
||||
torchmetrics
|
||||
torchtyping
|
||||
facenet_pytorch==2.5.0
|
||||
ffmpeg-python==0.2.0
|
||||
bitsandbytes==0.43.1
|
||||
decord==0.6.0
|
||||
diffusers==0.24.0
|
||||
einops==0.8.0
|
||||
insightface==0.7.3
|
||||
librosa==0.10.2.post1
|
||||
mediapipe[vision]==0.10.14
|
||||
mlflow==2.13.1
|
||||
moviepy==1.0.3
|
||||
numpy==1.26.4
|
||||
omegaconf==2.3.0
|
||||
onnx2torch==1.5.14
|
||||
onnx==1.16.1
|
||||
onnxruntime-gpu==1.18.0
|
||||
opencv-contrib-python==4.9.0.80
|
||||
opencv-python-headless==4.9.0.80
|
||||
opencv-python==4.9.0.80
|
||||
pillow==10.3.0
|
||||
setuptools==70.0.0
|
||||
tqdm==4.66.4
|
||||
transformers==4.39.2
|
||||
xformers==0.0.25.post1
|
||||
isort==5.13.2
|
||||
pylint==3.2.2
|
||||
pre-commit==3.7.1
|
||||
gradio==4.36.1
|
||||
huggingface-hub==0.25.2
|
||||
deepspeed==0.15.2
|
||||
tensorboard
|
||||
torchaudio==2.2.2
|
||||
denoiser==0.1.5
|
||||
soundfile==0.13.1
|
||||
deepface==0.0.93
|
||||
ultralytics==8.3.82
|
||||
tf-keras==2.18.0
|
||||
@@ -0,0 +1,114 @@
|
||||
# pylint: disable=R0801
|
||||
"""
|
||||
This module is used to extract meta information from video directories.
|
||||
|
||||
It takes in two command-line arguments: `root_path` and `dataset_name`. The `root_path`
|
||||
specifies the path to the video directory, while the `dataset_name` specifies the name
|
||||
of the dataset. The module then collects all the video folder paths, and for each video
|
||||
folder, it checks if a mask path and a face embedding path exist. If they do, it appends
|
||||
a dictionary containing the image path, mask path, and face embedding path to a list.
|
||||
|
||||
Finally, the module writes the list of dictionaries to a JSON file with the filename
|
||||
constructed using the `dataset_name`.
|
||||
|
||||
Usage:
|
||||
python tools/extract_meta_info_stage1.py --root_path /path/to/video_dir --dataset_name hdtf
|
||||
|
||||
"""
|
||||
|
||||
import argparse
|
||||
import json
|
||||
import os
|
||||
from pathlib import Path
|
||||
import random
|
||||
import torch
|
||||
import tqdm
|
||||
|
||||
|
||||
def collect_video_folder_paths(root_path: Path) -> list:
|
||||
"""
|
||||
Collect all video folder paths from the root path.
|
||||
|
||||
Args:
|
||||
root_path (Path): The root directory containing video folders.
|
||||
|
||||
Returns:
|
||||
list: List of video folder paths.
|
||||
"""
|
||||
return [frames_dir.resolve() for frames_dir in root_path.iterdir() if frames_dir.is_dir()]
|
||||
|
||||
|
||||
def construct_meta_info(frames_dir_path: Path) -> dict:
|
||||
"""
|
||||
Construct meta information for a given frames directory.
|
||||
|
||||
Args:
|
||||
frames_dir_path (Path): The path to the frames directory.
|
||||
|
||||
Returns:
|
||||
dict: A dictionary containing the meta information for the frames directory, or None if the required files do not exist.
|
||||
"""
|
||||
audio_emb_path = str(frames_dir_path).replace("images", "audio_emb_whisper") + ".pt"
|
||||
video_path = str(frames_dir_path).replace("images", "videos") + ".mp4"
|
||||
|
||||
if not (os.path.isfile(audio_emb_path) and os.path.isfile(video_path)):
|
||||
return None
|
||||
|
||||
# if torch.load(audio_emb_path) is None:
|
||||
# print(f"audio emb is None: {audio_emb_path}")
|
||||
# return None
|
||||
# if not os.path.exists(video_path):
|
||||
# print("no video exists")
|
||||
# return None
|
||||
|
||||
return {
|
||||
"image_path": str(frames_dir_path),
|
||||
"video_path": video_path,
|
||||
"audio_emb": audio_emb_path
|
||||
}
|
||||
|
||||
|
||||
def main():
|
||||
"""
|
||||
Main function to extract meta info for training.
|
||||
"""
|
||||
parser = argparse.ArgumentParser()
|
||||
parser.add_argument("-r", "--root_path", type=str,
|
||||
required=True, help="Root path of the video directories")
|
||||
parser.add_argument("-n", "--dataset_name", type=str,
|
||||
required=True, help="Name of the dataset")
|
||||
parser.add_argument("--meta_info_name", type=str,
|
||||
help="Name of the meta information file")
|
||||
|
||||
args = parser.parse_args()
|
||||
|
||||
if args.meta_info_name is None:
|
||||
args.meta_info_name = args.dataset_name
|
||||
|
||||
image_dir = Path(args.root_path) / "images"
|
||||
output_dir = Path("./data")
|
||||
output_dir.mkdir(exist_ok=True)
|
||||
|
||||
# Collect all video folder paths
|
||||
frames_dir_paths = collect_video_folder_paths(image_dir)
|
||||
random.shuffle(frames_dir_paths)
|
||||
|
||||
meta_infos = []
|
||||
for frames_dir_path in tqdm.tqdm(frames_dir_paths):
|
||||
meta_info = construct_meta_info(frames_dir_path)
|
||||
if meta_info:
|
||||
meta_infos.append(meta_info)
|
||||
|
||||
output_file = output_dir / f"{args.meta_info_name}_train.json"
|
||||
with output_file.open("w", encoding="utf-8") as f:
|
||||
json.dump(meta_infos[:-100], f, indent=4)
|
||||
|
||||
output_file = output_dir / f"{args.meta_info_name}_val.json"
|
||||
with output_file.open("w", encoding="utf-8") as f:
|
||||
json.dump(meta_infos[-100:], f, indent=4)
|
||||
|
||||
print(f"Final data count: {len(meta_infos)}")
|
||||
|
||||
|
||||
if __name__ == "__main__":
|
||||
main()
|
||||
@@ -0,0 +1,370 @@
|
||||
# pylint: disable=W0718
|
||||
"""
|
||||
This module is responsible for processing images, particularly for face-related tasks.
|
||||
It uses various libraries such as OpenCV, NumPy, and InsightFace to perform tasks like
|
||||
face detection, augmentation, and mask rendering. The ImageProcessor class encapsulates
|
||||
the functionality for these operations.
|
||||
"""
|
||||
import os
|
||||
from typing import List
|
||||
|
||||
import cv2
|
||||
import mediapipe as mp
|
||||
import numpy as np
|
||||
import torch
|
||||
from insightface.app import FaceAnalysis
|
||||
from PIL import Image, ImageOps
|
||||
from torchvision import transforms
|
||||
|
||||
from ..utils.util import (blur_mask, get_landmark_overframes, get_mask,
|
||||
get_union_face_mask, get_union_lip_mask)
|
||||
|
||||
MEAN = 0.5
|
||||
STD = 0.5
|
||||
|
||||
class ImageProcessor:
|
||||
"""
|
||||
ImageProcessor is a class responsible for processing images, particularly for face-related tasks.
|
||||
It takes in an image and performs various operations such as augmentation, face detection,
|
||||
face embedding extraction, and rendering a face mask. The processed images are then used for
|
||||
further analysis or recognition purposes.
|
||||
|
||||
Attributes:
|
||||
img_size (int): The size of the image to be processed.
|
||||
face_analysis_model_path (str): The path to the face analysis model.
|
||||
|
||||
Methods:
|
||||
preprocess(source_image_path, cache_dir):
|
||||
Preprocesses the input image by performing augmentation, face detection,
|
||||
face embedding extraction, and rendering a face mask.
|
||||
|
||||
close():
|
||||
Closes the ImageProcessor and releases any resources being used.
|
||||
|
||||
_augmentation(images, transform, state=None):
|
||||
Applies image augmentation to the input images using the given transform and state.
|
||||
|
||||
__enter__():
|
||||
Enters a runtime context and returns the ImageProcessor object.
|
||||
|
||||
__exit__(_exc_type, _exc_val, _exc_tb):
|
||||
Exits a runtime context and handles any exceptions that occurred during the processing.
|
||||
"""
|
||||
def __init__(self, img_size, face_analysis_model_path) -> None:
|
||||
self.img_size = img_size
|
||||
|
||||
self.pixel_transform = transforms.Compose(
|
||||
[
|
||||
transforms.Resize(self.img_size),
|
||||
transforms.ToTensor(),
|
||||
transforms.Normalize([MEAN], [STD]),
|
||||
]
|
||||
)
|
||||
|
||||
self.cond_transform = transforms.Compose(
|
||||
[
|
||||
transforms.Resize(self.img_size),
|
||||
transforms.ToTensor(),
|
||||
]
|
||||
)
|
||||
|
||||
self.attn_transform_64 = transforms.Compose(
|
||||
[
|
||||
transforms.Resize(
|
||||
(self.img_size[0] // 8, self.img_size[0] // 8)),
|
||||
transforms.ToTensor(),
|
||||
]
|
||||
)
|
||||
self.attn_transform_32 = transforms.Compose(
|
||||
[
|
||||
transforms.Resize(
|
||||
(self.img_size[0] // 16, self.img_size[0] // 16)),
|
||||
transforms.ToTensor(),
|
||||
]
|
||||
)
|
||||
self.attn_transform_16 = transforms.Compose(
|
||||
[
|
||||
transforms.Resize(
|
||||
(self.img_size[0] // 32, self.img_size[0] // 32)),
|
||||
transforms.ToTensor(),
|
||||
]
|
||||
)
|
||||
self.attn_transform_8 = transforms.Compose(
|
||||
[
|
||||
transforms.Resize(
|
||||
(self.img_size[0] // 64, self.img_size[0] // 64)),
|
||||
transforms.ToTensor(),
|
||||
]
|
||||
)
|
||||
|
||||
self.face_analysis = FaceAnalysis(
|
||||
name="",
|
||||
root=face_analysis_model_path,
|
||||
providers=["CUDAExecutionProvider", "CPUExecutionProvider"],
|
||||
)
|
||||
self.face_analysis.prepare(ctx_id=0, det_size=(640, 640))
|
||||
|
||||
def preprocess(self, frame):
|
||||
"""
|
||||
Apply preprocessing to the source image to prepare for face analysis.
|
||||
|
||||
Parameters:
|
||||
source_image_path (str): The path to the source image.
|
||||
cache_dir (str): The directory to cache intermediate results.
|
||||
|
||||
Returns:
|
||||
None
|
||||
"""
|
||||
frame_h, frame_w, _ = frame.shape
|
||||
# 2.1 detect face
|
||||
faces = self.face_analysis.get(frame)
|
||||
if not faces:
|
||||
print("No faces detected in the image. Using the entire image as the face region.")
|
||||
bbox = (0,0,0,0)
|
||||
face_emb = np.zeros(512)
|
||||
frame_crop = frame
|
||||
else:
|
||||
# Sort faces by size and select the largest one
|
||||
face = sorted(faces, key=lambda x: (x["bbox"][2] - x["bbox"][0]) * (x["bbox"][3] - x["bbox"][1]), reverse=True)[0]
|
||||
face_emb = face["embedding"]
|
||||
w, h = face["bbox"][2] - face["bbox"][0], face["bbox"][3] - face["bbox"][1]
|
||||
bbox = (int(max(face["bbox"][0] - 0.1*w, 0)), int(face["bbox"][1]),
|
||||
int(min(face["bbox"][2] + 0.1*w, frame_w)), int(min(face["bbox"][3] + 0.1*h, frame_h)))
|
||||
frame_crop = frame[bbox[1]:bbox[3], bbox[0]:bbox[2]]
|
||||
|
||||
cv2.imwrite("./x.jpg", frame_crop)
|
||||
frame_crop = cv2.resize(frame_crop, (512,512))
|
||||
|
||||
pixel_values_ref_img = self._augmentation(Image.fromarray(frame_crop).convert("RGB"), self.pixel_transform)
|
||||
|
||||
return pixel_values_ref_img, face_emb, bbox, frame_crop
|
||||
|
||||
|
||||
def preprocess_frames(self, frames):
|
||||
"""
|
||||
Apply preprocessing to the 14 source images to prepare for face analysis.
|
||||
|
||||
Parameters:
|
||||
source_image_path (str): The path to the source image.
|
||||
cache_dir (str): The directory to cache intermediate results.
|
||||
|
||||
Returns:
|
||||
None
|
||||
"""
|
||||
def find_max_h(bbox_list):
|
||||
|
||||
h_max = 0
|
||||
for bbox in bbox_list:
|
||||
h = bbox[3]
|
||||
if h>h_max:
|
||||
h_max = h
|
||||
|
||||
return h_max
|
||||
|
||||
pixel_values_ref_img_list, bbox_list, bbox_final_list, frame_crop_list = [], [], [], []
|
||||
|
||||
face_emb = np.zeros(512)
|
||||
|
||||
for frame in frames:
|
||||
|
||||
# 2.1 detect face
|
||||
faces = self.face_analysis.get(frame)
|
||||
if not faces:
|
||||
print("No faces detected in the image. Using the entire image as the face region.")
|
||||
bbox_list.append((0,0,0,0))
|
||||
|
||||
else:
|
||||
# Sort faces by size and select the largest one
|
||||
face = sorted(faces, key=lambda x: (x["bbox"][2] - x["bbox"][0]) * (x["bbox"][3] - x["bbox"][1]), reverse=True)[0]
|
||||
face_emb = face["embedding"]
|
||||
bbox_list.append(face["bbox"])
|
||||
|
||||
h_max = find_max_h(bbox_list)
|
||||
|
||||
for frame, bbox in zip(frames, bbox_list):
|
||||
|
||||
frame_h, frame_w, _ = frame.shape
|
||||
|
||||
w, h = bbox[2] - bbox[0], bbox[3] - bbox[1]
|
||||
|
||||
if h_max < 1.4*h:
|
||||
h = h_max
|
||||
|
||||
bbox_final = (int(max(bbox[0] - 0.1*w, 0)), int(bbox[1]),
|
||||
int(min(bbox[2] + 0.1*w, frame_w)), int(min(bbox[1] + 1.1*h, frame_h)))
|
||||
bbox_final_list.append(bbox_final)
|
||||
|
||||
frame_crop = frame[bbox_final[1]:bbox_final[3], bbox_final[0]:bbox_final[2]]
|
||||
frame_crop_list.append(cv2.resize(frame_crop, (512,512)))
|
||||
|
||||
pixel_values_ref_img_list.append(self._augmentation(Image.fromarray(frame_crop).convert("RGB"), self.pixel_transform).unsqueeze(0))
|
||||
|
||||
return pixel_values_ref_img_list, face_emb, bbox_list, frame_crop_list
|
||||
|
||||
|
||||
def close(self):
|
||||
"""
|
||||
Closes the ImageProcessor and releases any resources held by the FaceAnalysis instance.
|
||||
|
||||
Args:
|
||||
self: The ImageProcessor instance.
|
||||
|
||||
Returns:
|
||||
None.
|
||||
"""
|
||||
for _, model in self.face_analysis.models.items():
|
||||
if hasattr(model, "Dispose"):
|
||||
model.Dispose()
|
||||
|
||||
def _augmentation(self, images, transform, state=None):
|
||||
if state is not None:
|
||||
torch.set_rng_state(state)
|
||||
if isinstance(images, List):
|
||||
transformed_images = [transform(img) for img in images]
|
||||
ret_tensor = torch.stack(transformed_images, dim=0) # (f, c, h, w)
|
||||
else:
|
||||
ret_tensor = transform(images) # (c, h, w)
|
||||
return ret_tensor
|
||||
|
||||
def __enter__(self):
|
||||
return self
|
||||
|
||||
def __exit__(self, _exc_type, _exc_val, _exc_tb):
|
||||
self.close()
|
||||
|
||||
|
||||
class ImageProcessorForDataProcessing():
|
||||
"""
|
||||
ImageProcessor is a class responsible for processing images, particularly for face-related tasks.
|
||||
It takes in an image and performs various operations such as augmentation, face detection,
|
||||
face embedding extraction, and rendering a face mask. The processed images are then used for
|
||||
further analysis or recognition purposes.
|
||||
|
||||
Attributes:
|
||||
img_size (int): The size of the image to be processed.
|
||||
face_analysis_model_path (str): The path to the face analysis model.
|
||||
|
||||
Methods:
|
||||
preprocess(source_image_path, cache_dir):
|
||||
Preprocesses the input image by performing augmentation, face detection,
|
||||
face embedding extraction, and rendering a face mask.
|
||||
|
||||
close():
|
||||
Closes the ImageProcessor and releases any resources being used.
|
||||
|
||||
_augmentation(images, transform, state=None):
|
||||
Applies image augmentation to the input images using the given transform and state.
|
||||
|
||||
__enter__():
|
||||
Enters a runtime context and returns the ImageProcessor object.
|
||||
|
||||
__exit__(_exc_type, _exc_val, _exc_tb):
|
||||
Exits a runtime context and handles any exceptions that occurred during the processing.
|
||||
"""
|
||||
def __init__(self, face_analysis_model_path, landmark_model_path, step) -> None:
|
||||
if step == 2:
|
||||
self.face_analysis = FaceAnalysis(
|
||||
name="",
|
||||
root=face_analysis_model_path,
|
||||
providers=["CUDAExecutionProvider", "CPUExecutionProvider"],
|
||||
)
|
||||
self.face_analysis.prepare(ctx_id=0, det_size=(640, 640))
|
||||
self.landmarker = None
|
||||
else:
|
||||
BaseOptions = mp.tasks.BaseOptions
|
||||
FaceLandmarker = mp.tasks.vision.FaceLandmarker
|
||||
FaceLandmarkerOptions = mp.tasks.vision.FaceLandmarkerOptions
|
||||
VisionRunningMode = mp.tasks.vision.RunningMode
|
||||
# Create a face landmarker instance with the video mode:
|
||||
options = FaceLandmarkerOptions(
|
||||
base_options=BaseOptions(model_asset_path=landmark_model_path),
|
||||
running_mode=VisionRunningMode.IMAGE,
|
||||
)
|
||||
self.landmarker = None #FaceLandmarker.create_from_options(options)
|
||||
self.face_analysis = None
|
||||
|
||||
def preprocess(self, source_image_path: str):
|
||||
"""
|
||||
Apply preprocessing to the source image to prepare for face analysis.
|
||||
|
||||
Parameters:
|
||||
source_image_path (str): The path to the source image.
|
||||
cache_dir (str): The directory to cache intermediate results.
|
||||
|
||||
Returns:
|
||||
None
|
||||
"""
|
||||
# 1. get face embdeding
|
||||
face_mask, face_emb, sep_pose_mask, sep_face_mask, sep_lip_mask = None, None, None, None, None
|
||||
if self.face_analysis:
|
||||
frame = sorted(os.listdir(source_image_path))[0]
|
||||
source_image = Image.open(
|
||||
os.path.join(source_image_path, frame))
|
||||
ref_image_pil = source_image.convert("RGB")
|
||||
|
||||
padding = 250
|
||||
ref_image_pil = ImageOps.expand(ref_image_pil, border=padding, fill=0)
|
||||
|
||||
# 2.1 detect face
|
||||
faces = self.face_analysis.get(cv2.cvtColor(
|
||||
np.array(ref_image_pil.copy()), cv2.COLOR_RGB2BGR))
|
||||
|
||||
# use max size face
|
||||
face = sorted(faces, key=lambda x: (
|
||||
x["bbox"][2] - x["bbox"][0]) * (x["bbox"][3] - x["bbox"][1]))[-1]
|
||||
# 2.2 face embedding
|
||||
face_emb = face["embedding"]
|
||||
if face_emb is None:
|
||||
print(f"error in image embedding : {source_image_path}")
|
||||
raise
|
||||
|
||||
if self.landmarker:
|
||||
# 3.1 get landmark
|
||||
landmarks, height, width = get_landmark_overframes(
|
||||
self.landmarker, source_image_path)
|
||||
assert len(landmarks) == len(os.listdir(source_image_path))
|
||||
|
||||
# 3 render face and lip mask
|
||||
face_mask = get_union_face_mask(landmarks, height, width)
|
||||
lip_mask = get_union_lip_mask(landmarks, height, width)
|
||||
|
||||
# 4 gaussian blur
|
||||
blur_face_mask = blur_mask(face_mask, (64, 64), (51, 51))
|
||||
blur_lip_mask = blur_mask(lip_mask, (64, 64), (31, 31))
|
||||
|
||||
# 5 seperate mask
|
||||
sep_face_mask = cv2.subtract(blur_face_mask, blur_lip_mask)
|
||||
sep_pose_mask = 255.0 - blur_face_mask
|
||||
sep_lip_mask = blur_lip_mask
|
||||
|
||||
return face_mask, face_emb, sep_pose_mask, sep_face_mask, sep_lip_mask
|
||||
|
||||
def close(self):
|
||||
"""
|
||||
Closes the ImageProcessor and releases any resources held by the FaceAnalysis instance.
|
||||
|
||||
Args:
|
||||
self: The ImageProcessor instance.
|
||||
|
||||
Returns:
|
||||
None.
|
||||
"""
|
||||
for _, model in self.face_analysis.models.items():
|
||||
if hasattr(model, "Dispose"):
|
||||
model.Dispose()
|
||||
|
||||
def _augmentation(self, images, transform, state=None):
|
||||
if state is not None:
|
||||
torch.set_rng_state(state)
|
||||
if isinstance(images, List):
|
||||
transformed_images = [transform(img) for img in images]
|
||||
ret_tensor = torch.stack(transformed_images, dim=0) # (f, c, h, w)
|
||||
else:
|
||||
ret_tensor = transform(images) # (c, h, w)
|
||||
return ret_tensor
|
||||
|
||||
def __enter__(self):
|
||||
return self
|
||||
|
||||
def __exit__(self, _exc_type, _exc_val, _exc_tb):
|
||||
self.close()
|
||||
@@ -0,0 +1,291 @@
|
||||
# pylint: disable=R0801
|
||||
"""
|
||||
This module contains the code for a dataset class called FaceMaskDataset, which is used to process and
|
||||
load image data related to face masks. The dataset class inherits from the PyTorch Dataset class and
|
||||
provides methods for data augmentation, getting items from the dataset, and determining the length of the
|
||||
dataset. The module also includes imports for necessary libraries such as json, random, pathlib, torch,
|
||||
PIL, and transformers.
|
||||
"""
|
||||
|
||||
import json
|
||||
from pathlib import Path
|
||||
import cv2
|
||||
import numpy as np
|
||||
|
||||
import torch
|
||||
from PIL import Image
|
||||
from torch.utils.data import Dataset
|
||||
from torchvision import transforms
|
||||
from transformers import CLIPImageProcessor
|
||||
|
||||
|
||||
class FaceMaskDatasetVal(Dataset):
|
||||
"""
|
||||
FaceMaskDataset is a custom dataset for face mask images.
|
||||
|
||||
Args:
|
||||
img_size (int): The size of the input images.
|
||||
drop_ratio (float, optional): The ratio of dropped pixels during data augmentation. Defaults to 0.1.
|
||||
data_meta_paths (list, optional): The paths to the metadata files containing image paths and labels. Defaults to ["./data/HDTF_meta.json"].
|
||||
sample_margin (int, optional): The margin for sampling regions in the image. Defaults to 30.
|
||||
|
||||
Attributes:
|
||||
img_size (int): The size of the input images.
|
||||
drop_ratio (float): The ratio of dropped pixels during data augmentation.
|
||||
data_meta_paths (list): The paths to the metadata files containing image paths and labels.
|
||||
sample_margin (int): The margin for sampling regions in the image.
|
||||
processor (CLIPImageProcessor): The image processor for preprocessing images.
|
||||
transform (transforms.Compose): The image augmentation transform.
|
||||
"""
|
||||
|
||||
def __init__(
|
||||
self,
|
||||
img_size,
|
||||
mask_path,
|
||||
drop_ratio=0.1,
|
||||
data_meta_path=None,
|
||||
sample_margin=30,
|
||||
):
|
||||
super().__init__()
|
||||
|
||||
self.img_size = img_size
|
||||
self.sample_margin = sample_margin
|
||||
|
||||
vid_meta = []
|
||||
with open(data_meta_path, "r", encoding="utf-8") as f:
|
||||
vid_meta.extend(json.load(f))
|
||||
self.vid_meta = vid_meta
|
||||
self.length = len(self.vid_meta)
|
||||
|
||||
self.mask = self.create_mask(mask_path)
|
||||
|
||||
self.clip_image_processor = CLIPImageProcessor()
|
||||
|
||||
self.transform = transforms.Compose(
|
||||
[
|
||||
transforms.Resize(self.img_size),
|
||||
transforms.ToTensor(),
|
||||
transforms.Normalize([0.5], [0.5]),
|
||||
]
|
||||
)
|
||||
|
||||
self.cond_transform = transforms.Compose(
|
||||
[
|
||||
transforms.Resize(self.img_size),
|
||||
transforms.ToTensor(),
|
||||
]
|
||||
)
|
||||
|
||||
self.drop_ratio = drop_ratio
|
||||
|
||||
def create_mask(self, mask_path):
|
||||
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 augmentation(self, image, transform, state=None):
|
||||
"""
|
||||
Apply data augmentation to the input image.
|
||||
|
||||
Args:
|
||||
image (PIL.Image): The input image.
|
||||
transform (torchvision.transforms.Compose): The data augmentation transforms.
|
||||
state (dict, optional): The random state for reproducibility. Defaults to None.
|
||||
|
||||
Returns:
|
||||
PIL.Image: The augmented image.
|
||||
"""
|
||||
if state is not None:
|
||||
torch.set_rng_state(state)
|
||||
return transform(image)
|
||||
|
||||
def __getitem__(self, index):
|
||||
video_meta = self.vid_meta[index]
|
||||
video_path = video_meta["image_path"]
|
||||
|
||||
video_frames = sorted(Path(video_path).iterdir())
|
||||
video_length = len(video_frames)
|
||||
|
||||
margin = min(self.sample_margin, video_length)
|
||||
|
||||
ref_img_idx = np.random.randint(0, video_length - 1)
|
||||
if ref_img_idx + margin < (video_length - 1):
|
||||
tgt_img_idx = np.random.randint(
|
||||
ref_img_idx + margin, video_length - 1)
|
||||
elif ref_img_idx - margin > 0:
|
||||
tgt_img_idx = np.random.randint(0, ref_img_idx - margin)
|
||||
else:
|
||||
tgt_img_idx = np.random.randint(0, video_length - 1)
|
||||
|
||||
ref_img_pil = Image.open(video_frames[ref_img_idx])
|
||||
tgt_img_pil = Image.open(video_frames[tgt_img_idx])
|
||||
|
||||
assert ref_img_pil is not None, "Fail to load reference image."
|
||||
assert tgt_img_pil is not None, "Fail to load target image."
|
||||
|
||||
audio_emb_path = video_meta["audio_emb"]
|
||||
audio_emb = torch.load(audio_emb_path)
|
||||
audio_tensor = audio_emb[tgt_img_idx]
|
||||
|
||||
sample = {
|
||||
"video_dir": video_path,
|
||||
"img": tgt_img_pil,
|
||||
"ref_img": ref_img_pil,
|
||||
"audio_tensor": audio_tensor,
|
||||
"mask":self.mask,
|
||||
}
|
||||
|
||||
return sample
|
||||
|
||||
def __len__(self):
|
||||
return len(self.vid_meta)
|
||||
|
||||
|
||||
|
||||
class FaceMaskDataset(Dataset):
|
||||
"""
|
||||
FaceMaskDataset is a custom dataset for face mask images.
|
||||
|
||||
Args:
|
||||
img_size (int): The size of the input images.
|
||||
drop_ratio (float, optional): The ratio of dropped pixels during data augmentation. Defaults to 0.1.
|
||||
data_meta_paths (list, optional): The paths to the metadata files containing image paths and labels. Defaults to ["./data/HDTF_meta.json"].
|
||||
sample_margin (int, optional): The margin for sampling regions in the image. Defaults to 30.
|
||||
|
||||
Attributes:
|
||||
img_size (int): The size of the input images.
|
||||
drop_ratio (float): The ratio of dropped pixels during data augmentation.
|
||||
data_meta_paths (list): The paths to the metadata files containing image paths and labels.
|
||||
sample_margin (int): The margin for sampling regions in the image.
|
||||
processor (CLIPImageProcessor): The image processor for preprocessing images.
|
||||
transform (transforms.Compose): The image augmentation transform.
|
||||
"""
|
||||
|
||||
def __init__(
|
||||
self,
|
||||
img_size,
|
||||
mask_path,
|
||||
drop_ratio=0.1,
|
||||
data_meta_path=None,
|
||||
sample_margin=30,
|
||||
):
|
||||
super().__init__()
|
||||
|
||||
self.img_size = img_size
|
||||
self.sample_margin = sample_margin
|
||||
|
||||
vid_meta = []
|
||||
with open(data_meta_path, "r", encoding="utf-8") as f:
|
||||
vid_meta.extend(json.load(f))
|
||||
self.vid_meta = vid_meta
|
||||
self.length = len(self.vid_meta)
|
||||
|
||||
self.mask = self.create_mask(mask_path)
|
||||
|
||||
self.clip_image_processor = CLIPImageProcessor()
|
||||
|
||||
self.transform = transforms.Compose(
|
||||
[
|
||||
transforms.Resize(self.img_size),
|
||||
transforms.ToTensor(),
|
||||
transforms.Normalize([0.5], [0.5]),
|
||||
]
|
||||
)
|
||||
|
||||
self.cond_transform = transforms.Compose(
|
||||
[
|
||||
transforms.Resize(self.img_size),
|
||||
transforms.ToTensor(),
|
||||
]
|
||||
)
|
||||
|
||||
self.drop_ratio = drop_ratio
|
||||
|
||||
def create_mask(self, mask_path):
|
||||
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 augmentation(self, image, transform, state=None):
|
||||
"""
|
||||
Apply data augmentation to the input image.
|
||||
|
||||
Args:
|
||||
image (PIL.Image): The input image.
|
||||
transform (torchvision.transforms.Compose): The data augmentation transforms.
|
||||
state (dict, optional): The random state for reproducibility. Defaults to None.
|
||||
|
||||
Returns:
|
||||
PIL.Image: The augmented image.
|
||||
"""
|
||||
if state is not None:
|
||||
torch.set_rng_state(state)
|
||||
return transform(image)
|
||||
|
||||
def __getitem__(self, index):
|
||||
video_meta = self.vid_meta[index]
|
||||
video_path = video_meta["image_path"]
|
||||
|
||||
video_frames = sorted(Path(video_path).iterdir())
|
||||
video_length = len(video_frames)
|
||||
|
||||
margin = min(self.sample_margin, video_length)
|
||||
|
||||
ref_img_idx = np.random.randint(0, video_length - 1)
|
||||
if ref_img_idx + margin < (video_length -1) :
|
||||
tgt_img_idx = np.random.randint(
|
||||
ref_img_idx + margin, video_length - 1)
|
||||
elif ref_img_idx - margin > 0:
|
||||
tgt_img_idx = np.random.randint(0, ref_img_idx - margin)
|
||||
else:
|
||||
tgt_img_idx = np.random.randint(0, video_length - 1)
|
||||
|
||||
ref_img_pil = Image.open(video_frames[ref_img_idx])
|
||||
tgt_img_pil = Image.open(video_frames[tgt_img_idx])
|
||||
|
||||
assert ref_img_pil is not None, "Fail to load reference image."
|
||||
assert tgt_img_pil is not None, "Fail to load target image."
|
||||
|
||||
state = torch.get_rng_state()
|
||||
tgt_img = self.augmentation(tgt_img_pil, self.transform, state)
|
||||
ref_img_vae = self.augmentation(
|
||||
ref_img_pil, self.transform, state)
|
||||
|
||||
audio_emb_path = video_meta["audio_emb"]
|
||||
audio_emb = torch.load(audio_emb_path)
|
||||
audio_tensor = audio_emb[tgt_img_idx]
|
||||
|
||||
sample = {
|
||||
"video_dir": video_path,
|
||||
"img": tgt_img,
|
||||
"ref_img": ref_img_vae,
|
||||
"audio_tensor": audio_tensor,
|
||||
"mask":self.mask,
|
||||
}
|
||||
|
||||
return sample
|
||||
|
||||
def __len__(self):
|
||||
return len(self.vid_meta)
|
||||
|
||||
|
||||
if __name__ == "__main__":
|
||||
data = FaceMaskDataset(img_size=(512, 512))
|
||||
train_dataloader = torch.utils.data.DataLoader(
|
||||
data, batch_size=4, shuffle=True, num_workers=1
|
||||
)
|
||||
for step, batch in enumerate(train_dataloader):
|
||||
print(batch["tgt_mask"].shape)
|
||||
break
|
||||
@@ -0,0 +1,456 @@
|
||||
# pylint: disable=R0801
|
||||
"""
|
||||
talking_video_dataset.py
|
||||
|
||||
This module defines the TalkingVideoDataset class, a custom PyTorch dataset
|
||||
for handling talking video data. The dataset uses video files, masks, and
|
||||
embeddings to prepare data for tasks such as video generation and
|
||||
speech-driven video animation.
|
||||
|
||||
Classes:
|
||||
TalkingVideoDataset
|
||||
|
||||
Dependencies:
|
||||
json
|
||||
random
|
||||
torch
|
||||
decord.VideoReader, decord.cpu
|
||||
PIL.Image
|
||||
torch.utils.data.Dataset
|
||||
torchvision.transforms
|
||||
|
||||
Example:
|
||||
from talking_video_dataset import TalkingVideoDataset
|
||||
from torch.utils.data import DataLoader
|
||||
|
||||
# Example configuration for the Wav2Vec model
|
||||
class Wav2VecConfig:
|
||||
def __init__(self, audio_type, model_scale, features):
|
||||
self.audio_type = audio_type
|
||||
self.model_scale = model_scale
|
||||
self.features = features
|
||||
|
||||
wav2vec_cfg = Wav2VecConfig(audio_type="wav2vec2", model_scale="base", features="feature")
|
||||
|
||||
# Initialize dataset
|
||||
dataset = TalkingVideoDataset(
|
||||
img_size=(512, 512),
|
||||
sample_rate=16000,
|
||||
audio_margin=2,
|
||||
n_sample_frames=16,
|
||||
data_meta_paths=["path/to/meta1.json", "path/to/meta2.json"],
|
||||
wav2vec_cfg=wav2vec_cfg,
|
||||
)
|
||||
|
||||
# Initialize dataloader
|
||||
dataloader = DataLoader(dataset, batch_size=4, shuffle=True)
|
||||
|
||||
# Fetch one batch of data
|
||||
batch = next(iter(dataloader))
|
||||
print(batch["pixel_values_vid"].shape) # Example output: (4, 16, 3, 512, 512)
|
||||
|
||||
The TalkingVideoDataset class provides methods for loading video frames, masks,
|
||||
audio embeddings, and other relevant data, applying transformations, and preparing
|
||||
the data for training and evaluation in a deep learning pipeline.
|
||||
|
||||
Attributes:
|
||||
img_size (tuple): The dimensions to resize the video frames to.
|
||||
sample_rate (int): The audio sample rate.
|
||||
audio_margin (int): The margin for audio sampling.
|
||||
n_sample_frames (int): The number of sample frames.
|
||||
data_meta_paths (list): List of paths to the JSON metadata files.
|
||||
wav2vec_cfg (object): Configuration for the Wav2Vec model.
|
||||
|
||||
Methods:
|
||||
augmentation(images, transform, state=None): Apply transformation to input images.
|
||||
__getitem__(index): Get a sample from the dataset at the specified index.
|
||||
__len__(): Return the length of the dataset.
|
||||
"""
|
||||
|
||||
import json
|
||||
from typing import List
|
||||
import cv2
|
||||
import numpy as np
|
||||
|
||||
import torch
|
||||
from decord import VideoReader, cpu
|
||||
from PIL import Image
|
||||
from torch.utils.data import Dataset
|
||||
from torchvision import transforms
|
||||
|
||||
class TalkingVideoDatasetVal(Dataset):
|
||||
"""
|
||||
A dataset class for processing talking video data.
|
||||
|
||||
Args:
|
||||
img_size (tuple, optional): The size of the output images. Defaults to (512, 512).
|
||||
sample_rate (int, optional): The sample rate of the audio data. Defaults to 16000.
|
||||
audio_margin (int, optional): The margin for the audio data. Defaults to 2.
|
||||
n_sample_frames (int, optional): The number of sample frames. Defaults to 16.
|
||||
data_meta_paths (list, optional): The paths to the data metadata. Defaults to None.
|
||||
wav2vec_cfg (dict, optional): The configuration for the wav2vec model. Defaults to None.
|
||||
|
||||
Attributes:
|
||||
img_size (tuple): The size of the output images.
|
||||
sample_rate (int): The sample rate of the audio data.
|
||||
audio_margin (int): The margin for the audio data.
|
||||
n_sample_frames (int): The number of sample frames.
|
||||
data_meta_paths (list): The paths to the data metadata.
|
||||
wav2vec_cfg (dict): The configuration for the wav2vec model.
|
||||
"""
|
||||
|
||||
def __init__(
|
||||
self,
|
||||
img_size=(512, 512),
|
||||
mask_path=None,
|
||||
sample_rate=16000,
|
||||
audio_margin=2,
|
||||
n_sample_frames=16,
|
||||
data_meta_path=None,
|
||||
wav2vec_cfg=None,
|
||||
):
|
||||
super().__init__()
|
||||
self.sample_rate = sample_rate
|
||||
self.img_size = img_size
|
||||
self.audio_margin = audio_margin
|
||||
self.n_sample_frames = n_sample_frames
|
||||
self.audio_type = wav2vec_cfg.audio_type
|
||||
self.audio_model = wav2vec_cfg.model_scale
|
||||
self.audio_features = wav2vec_cfg.features
|
||||
|
||||
vid_meta = []
|
||||
with open(data_meta_path, "r", encoding="utf-8") as f:
|
||||
vid_meta.extend(json.load(f))
|
||||
self.vid_meta = vid_meta
|
||||
self.length = len(self.vid_meta)
|
||||
|
||||
self.mask_path = mask_path
|
||||
self.mask = self.create_mask(mask_path)
|
||||
|
||||
self.pixel_transform = transforms.Compose(
|
||||
[
|
||||
transforms.Resize(self.img_size),
|
||||
transforms.ToTensor(),
|
||||
transforms.Normalize([0.5], [0.5]),
|
||||
]
|
||||
)
|
||||
|
||||
self.cond_transform = transforms.Compose(
|
||||
[
|
||||
transforms.Resize(self.img_size),
|
||||
transforms.ToTensor(),
|
||||
]
|
||||
)
|
||||
self.attn_transform_64 = transforms.Compose(
|
||||
[
|
||||
transforms.Resize(
|
||||
(self.img_size[0] // 8, self.img_size[0] // 8)),
|
||||
transforms.ToTensor(),
|
||||
]
|
||||
)
|
||||
self.attn_transform_32 = transforms.Compose(
|
||||
[
|
||||
transforms.Resize(
|
||||
(self.img_size[0] // 16, self.img_size[0] // 16)),
|
||||
transforms.ToTensor(),
|
||||
]
|
||||
)
|
||||
self.attn_transform_16 = transforms.Compose(
|
||||
[
|
||||
transforms.Resize(
|
||||
(self.img_size[0] // 32, self.img_size[0] // 32)),
|
||||
transforms.ToTensor(),
|
||||
]
|
||||
)
|
||||
self.attn_transform_8 = transforms.Compose(
|
||||
[
|
||||
transforms.Resize(
|
||||
(self.img_size[0] // 64, self.img_size[0] // 64)),
|
||||
transforms.ToTensor(),
|
||||
]
|
||||
)
|
||||
|
||||
def create_mask(self, mask_path):
|
||||
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)
|
||||
# mask = mask.repeat(1,self.n_sample_frames,1,1)
|
||||
|
||||
return mask
|
||||
|
||||
def augmentation(self, images, transform, state=None):
|
||||
"""
|
||||
Apply the given transformation to the input images.
|
||||
|
||||
Args:
|
||||
images (List[PIL.Image] or PIL.Image): The input images to be transformed.
|
||||
transform (torchvision.transforms.Compose): The transformation to be applied to the images.
|
||||
state (torch.ByteTensor, optional): The state of the random number generator.
|
||||
If provided, it will set the RNG state to this value before applying the transformation. Defaults to None.
|
||||
|
||||
Returns:
|
||||
torch.Tensor: The transformed images as a tensor.
|
||||
If the input was a list of images, the tensor will have shape (f, c, h, w),
|
||||
where f is the number of images, c is the number of channels, h is the height, and w is the width.
|
||||
If the input was a single image, the tensor will have shape (c, h, w),
|
||||
where c is the number of channels, h is the height, and w is the width.
|
||||
"""
|
||||
if state is not None:
|
||||
torch.set_rng_state(state)
|
||||
if isinstance(images, List):
|
||||
transformed_images = [transform(img) for img in images]
|
||||
ret_tensor = torch.stack(transformed_images, dim=0) # (f, c, h, w)
|
||||
else:
|
||||
ret_tensor = transform(images) # (c, h, w)
|
||||
return ret_tensor
|
||||
|
||||
def __getitem__(self, index):
|
||||
video_meta = self.vid_meta[index]
|
||||
video_path = video_meta["video_path"] #"/workspace/preprocess/videos/ga3GzQWqWWg_C0.mp4"
|
||||
same = False #random.choice((True, False))
|
||||
if same:
|
||||
audio_emb_path = video_meta["audio_emb"]
|
||||
audio_path = audio_emb_path.replace("audio_emb_whisper", "audios")[:-2] + "wav"
|
||||
else:
|
||||
if (index+1) > (len(self.vid_meta)-1):
|
||||
new_index = index - 1
|
||||
else:
|
||||
new_index = index + 1
|
||||
video_meta_next = self.vid_meta[new_index]
|
||||
audio_emb_path = video_meta_next["audio_emb"]
|
||||
audio_path = audio_emb_path.replace("audio_emb_whisper", "audios")[:-2] + "wav"
|
||||
|
||||
# video_frames = VideoReader(video_path, ctx=cpu(0))
|
||||
video_frames = []
|
||||
video_stream = cv2.VideoCapture(video_path)
|
||||
while True:
|
||||
still_reading, frame = video_stream.read()
|
||||
if (not still_reading):
|
||||
video_stream.release()
|
||||
break
|
||||
frame = cv2.cvtColor(frame, cv2.COLOR_BGR2RGB)
|
||||
video_frames.append(frame)
|
||||
|
||||
frame_list = [
|
||||
Image.fromarray(video).convert("RGB") for video in video_frames[:5*25]] #.asnumpy()
|
||||
|
||||
audio_emb = torch.tensor(torch.load(audio_emb_path))
|
||||
|
||||
# transform
|
||||
state = torch.get_rng_state()
|
||||
pixel_values_vid = self.augmentation(frame_list, self.pixel_transform, state)
|
||||
|
||||
pixel_values_ref_img = pixel_values_vid[0]
|
||||
pixel_values_ref_img = pixel_values_ref_img.unsqueeze(0)
|
||||
|
||||
sample = {
|
||||
"video_dir": video_path,
|
||||
"pixel_values_vid": pixel_values_vid,
|
||||
"mask": self.mask,
|
||||
"audio_tensor": audio_emb,
|
||||
"pixel_values_ref_img": pixel_values_ref_img,
|
||||
"same" : same,
|
||||
"audio_path" : audio_path
|
||||
}
|
||||
|
||||
return sample
|
||||
|
||||
def __len__(self):
|
||||
return len(self.vid_meta)
|
||||
|
||||
|
||||
|
||||
class TalkingVideoDataset(Dataset):
|
||||
"""
|
||||
A dataset class for processing talking video data.
|
||||
|
||||
Args:
|
||||
img_size (tuple, optional): The size of the output images. Defaults to (512, 512).
|
||||
sample_rate (int, optional): The sample rate of the audio data. Defaults to 16000.
|
||||
audio_margin (int, optional): The margin for the audio data. Defaults to 2.
|
||||
n_sample_frames (int, optional): The number of sample frames. Defaults to 16.
|
||||
data_meta_paths (list, optional): The paths to the data metadata. Defaults to None.
|
||||
wav2vec_cfg (dict, optional): The configuration for the wav2vec model. Defaults to None.
|
||||
|
||||
Attributes:
|
||||
img_size (tuple): The size of the output images.
|
||||
sample_rate (int): The sample rate of the audio data.
|
||||
audio_margin (int): The margin for the audio data.
|
||||
n_sample_frames (int): The number of sample frames.
|
||||
data_meta_paths (list): The paths to the data metadata.
|
||||
wav2vec_cfg (dict): The configuration for the wav2vec model.
|
||||
"""
|
||||
|
||||
def __init__(
|
||||
self,
|
||||
img_size=(512, 512),
|
||||
mask_path=None,
|
||||
sample_rate=16000,
|
||||
audio_margin=2,
|
||||
n_sample_frames=16,
|
||||
data_meta_path=None,
|
||||
wav2vec_cfg=None,
|
||||
):
|
||||
super().__init__()
|
||||
self.sample_rate = sample_rate
|
||||
self.img_size = img_size
|
||||
self.audio_margin = audio_margin
|
||||
self.n_sample_frames = n_sample_frames
|
||||
self.audio_type = wav2vec_cfg.audio_type
|
||||
self.audio_model = wav2vec_cfg.model_scale
|
||||
self.audio_features = wav2vec_cfg.features
|
||||
|
||||
vid_meta = []
|
||||
with open(data_meta_path, "r", encoding="utf-8") as f:
|
||||
vid_meta.extend(json.load(f))
|
||||
self.vid_meta = vid_meta
|
||||
self.length = len(self.vid_meta)
|
||||
|
||||
self.mask_path = mask_path
|
||||
self.mask = self.create_mask(mask_path)
|
||||
|
||||
self.pixel_transform = transforms.Compose(
|
||||
[
|
||||
transforms.Resize(self.img_size),
|
||||
transforms.ToTensor(),
|
||||
transforms.Normalize([0.5], [0.5]),
|
||||
]
|
||||
)
|
||||
|
||||
self.cond_transform = transforms.Compose(
|
||||
[
|
||||
transforms.Resize(self.img_size),
|
||||
transforms.ToTensor(),
|
||||
]
|
||||
)
|
||||
self.attn_transform_64 = transforms.Compose(
|
||||
[
|
||||
transforms.Resize(
|
||||
(self.img_size[0] // 8, self.img_size[0] // 8)),
|
||||
transforms.ToTensor(),
|
||||
]
|
||||
)
|
||||
self.attn_transform_32 = transforms.Compose(
|
||||
[
|
||||
transforms.Resize(
|
||||
(self.img_size[0] // 16, self.img_size[0] // 16)),
|
||||
transforms.ToTensor(),
|
||||
]
|
||||
)
|
||||
self.attn_transform_16 = transforms.Compose(
|
||||
[
|
||||
transforms.Resize(
|
||||
(self.img_size[0] // 32, self.img_size[0] // 32)),
|
||||
transforms.ToTensor(),
|
||||
]
|
||||
)
|
||||
self.attn_transform_8 = transforms.Compose(
|
||||
[
|
||||
transforms.Resize(
|
||||
(self.img_size[0] // 64, self.img_size[0] // 64)),
|
||||
transforms.ToTensor(),
|
||||
]
|
||||
)
|
||||
|
||||
def create_mask(self, mask_path):
|
||||
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)
|
||||
mask = mask.repeat(1,self.n_sample_frames,1,1)
|
||||
|
||||
return mask
|
||||
|
||||
def augmentation(self, images, transform, state=None):
|
||||
"""
|
||||
Apply the given transformation to the input images.
|
||||
|
||||
Args:
|
||||
images (List[PIL.Image] or PIL.Image): The input images to be transformed.
|
||||
transform (torchvision.transforms.Compose): The transformation to be applied to the images.
|
||||
state (torch.ByteTensor, optional): The state of the random number generator.
|
||||
If provided, it will set the RNG state to this value before applying the transformation. Defaults to None.
|
||||
|
||||
Returns:
|
||||
torch.Tensor: The transformed images as a tensor.
|
||||
If the input was a list of images, the tensor will have shape (f, c, h, w),
|
||||
where f is the number of images, c is the number of channels, h is the height, and w is the width.
|
||||
If the input was a single image, the tensor will have shape (c, h, w),
|
||||
where c is the number of channels, h is the height, and w is the width.
|
||||
"""
|
||||
if state is not None:
|
||||
torch.set_rng_state(state)
|
||||
if isinstance(images, List):
|
||||
transformed_images = [transform(img) for img in images]
|
||||
ret_tensor = torch.stack(transformed_images, dim=0) # (f, c, h, w)
|
||||
else:
|
||||
ret_tensor = transform(images) # (c, h, w)
|
||||
return ret_tensor
|
||||
|
||||
def __getitem__(self, index):
|
||||
video_meta = self.vid_meta[index]
|
||||
video_path = video_meta["video_path"]
|
||||
audio_emb_path = video_meta["audio_emb"]
|
||||
|
||||
video_frames = VideoReader(video_path, ctx=cpu(0))
|
||||
assert (video_frames is not None and len(video_frames) > 0), "Fail to load video frames."
|
||||
audio_emb = torch.load(audio_emb_path)
|
||||
video_length = min(len(video_frames), audio_emb.shape[0])
|
||||
|
||||
assert (
|
||||
video_length
|
||||
> self.n_sample_frames + 2 * self.audio_margin
|
||||
)
|
||||
start_idx = np.random.randint(
|
||||
3,
|
||||
video_length - self.n_sample_frames - self.audio_margin - 1,
|
||||
)
|
||||
|
||||
videos = video_frames[start_idx : start_idx + self.n_sample_frames]
|
||||
|
||||
frame_list = [
|
||||
Image.fromarray(video).convert("RGB") for video in videos.asnumpy()
|
||||
]
|
||||
|
||||
indices = (
|
||||
torch.arange(2 * self.audio_margin + 1) - self.audio_margin
|
||||
) # Generates [-2, -1, 0, 1, 2]
|
||||
center_indices = torch.arange(
|
||||
start_idx,
|
||||
start_idx + self.n_sample_frames,
|
||||
)
|
||||
audio_tensor = audio_emb[center_indices]
|
||||
|
||||
ref_img_idx = np.random.randint(
|
||||
3,
|
||||
video_length - self.n_sample_frames - self.audio_margin - 1,
|
||||
)
|
||||
ref_img = video_frames[ref_img_idx].asnumpy()
|
||||
ref_img = Image.fromarray(ref_img)
|
||||
|
||||
# transform
|
||||
state = torch.get_rng_state()
|
||||
pixel_values_vid = self.augmentation(frame_list, self.pixel_transform, state)
|
||||
|
||||
pixel_values_ref_img = self.augmentation(ref_img, self.pixel_transform, state)
|
||||
pixel_values_ref_img = pixel_values_ref_img.unsqueeze(0)
|
||||
|
||||
sample = {
|
||||
"video_dir": video_path,
|
||||
"pixel_values_vid": pixel_values_vid,
|
||||
"mask": self.mask,
|
||||
"audio_tensor": audio_tensor,
|
||||
"pixel_values_ref_img": pixel_values_ref_img,
|
||||
}
|
||||
|
||||
return sample
|
||||
|
||||
def __len__(self):
|
||||
return len(self.vid_meta)
|
||||
@@ -0,0 +1,338 @@
|
||||
# Adapted from https://github.com/huggingface/diffusers/blob/main/src/diffusers/models/attention.py
|
||||
|
||||
from typing import Any, Dict, Optional
|
||||
|
||||
import torch
|
||||
from diffusers.models.attention import AdaLayerNorm, Attention, FeedForward
|
||||
from diffusers.models.embeddings import SinusoidalPositionalEmbedding
|
||||
from einops import rearrange
|
||||
from torch import nn
|
||||
|
||||
|
||||
class BasicTransformerBlock(nn.Module):
|
||||
r"""
|
||||
A basic Transformer block.
|
||||
|
||||
Parameters:
|
||||
dim (`int`): The number of channels in the input and output.
|
||||
num_attention_heads (`int`): The number of heads to use for multi-head attention.
|
||||
attention_head_dim (`int`): The number of channels in each head.
|
||||
dropout (`float`, *optional*, defaults to 0.0): The dropout probability to use.
|
||||
cross_attention_dim (`int`, *optional*): The size of the encoder_hidden_states vector for cross attention.
|
||||
activation_fn (`str`, *optional*, defaults to `"geglu"`): Activation function to be used in feed-forward.
|
||||
num_embeds_ada_norm (:
|
||||
obj: `int`, *optional*): The number of diffusion steps used during training. See `Transformer2DModel`.
|
||||
attention_bias (:
|
||||
obj: `bool`, *optional*, defaults to `False`): Configure if the attentions should contain a bias parameter.
|
||||
only_cross_attention (`bool`, *optional*):
|
||||
Whether to use only cross-attention layers. In this case two cross attention layers are used.
|
||||
double_self_attention (`bool`, *optional*):
|
||||
Whether to use two self-attention layers. In this case no cross attention layers are used.
|
||||
upcast_attention (`bool`, *optional*):
|
||||
Whether to upcast the attention computation to float32. This is useful for mixed precision training.
|
||||
norm_elementwise_affine (`bool`, *optional*, defaults to `True`):
|
||||
Whether to use learnable elementwise affine parameters for normalization.
|
||||
norm_type (`str`, *optional*, defaults to `"layer_norm"`):
|
||||
The normalization layer to use. Can be `"layer_norm"`, `"ada_norm"` or `"ada_norm_zero"`.
|
||||
final_dropout (`bool` *optional*, defaults to False):
|
||||
Whether to apply a final dropout after the last feed-forward layer.
|
||||
attention_type (`str`, *optional*, defaults to `"default"`):
|
||||
The type of attention to use. Can be `"default"` or `"gated"` or `"gated-text-image"`.
|
||||
positional_embeddings (`str`, *optional*, defaults to `None`):
|
||||
The type of positional embeddings to apply to.
|
||||
num_positional_embeddings (`int`, *optional*, defaults to `None`):
|
||||
The maximum number of positional embeddings to apply.
|
||||
"""
|
||||
|
||||
def __init__(
|
||||
self,
|
||||
dim: int,
|
||||
num_attention_heads: int,
|
||||
attention_head_dim: int,
|
||||
dropout=0.0,
|
||||
cross_attention_dim: Optional[int] = None,
|
||||
activation_fn: str = "geglu",
|
||||
num_embeds_ada_norm: Optional[int] = None,
|
||||
attention_bias: bool = False,
|
||||
only_cross_attention: bool = False,
|
||||
double_self_attention: bool = False,
|
||||
upcast_attention: bool = False,
|
||||
norm_elementwise_affine: bool = True,
|
||||
norm_type: str = "layer_norm", # 'layer_norm', 'ada_norm', 'ada_norm_zero', 'ada_norm_single'
|
||||
norm_eps: float = 1e-5,
|
||||
final_dropout: bool = False,
|
||||
attention_type: str = "default",
|
||||
positional_embeddings: Optional[str] = None,
|
||||
num_positional_embeddings: Optional[int] = None,
|
||||
):
|
||||
super().__init__()
|
||||
self.only_cross_attention = only_cross_attention
|
||||
|
||||
self.use_ada_layer_norm_zero = (
|
||||
num_embeds_ada_norm is not None
|
||||
) and norm_type == "ada_norm_zero"
|
||||
self.use_ada_layer_norm = (
|
||||
num_embeds_ada_norm is not None
|
||||
) and norm_type == "ada_norm"
|
||||
self.use_ada_layer_norm_single = norm_type == "ada_norm_single"
|
||||
self.use_layer_norm = norm_type == "layer_norm"
|
||||
|
||||
if norm_type in ("ada_norm", "ada_norm_zero") and num_embeds_ada_norm is None:
|
||||
raise ValueError(
|
||||
f"`norm_type` is set to {norm_type}, but `num_embeds_ada_norm` is not defined. Please make sure to"
|
||||
f" define `num_embeds_ada_norm` if setting `norm_type` to {norm_type}."
|
||||
)
|
||||
|
||||
if positional_embeddings and (num_positional_embeddings is None):
|
||||
raise ValueError(
|
||||
"If `positional_embedding` type is defined, `num_positition_embeddings` must also be defined."
|
||||
)
|
||||
|
||||
if positional_embeddings == "sinusoidal":
|
||||
self.pos_embed = SinusoidalPositionalEmbedding(
|
||||
dim, max_seq_length=num_positional_embeddings
|
||||
)
|
||||
else:
|
||||
self.pos_embed = None
|
||||
|
||||
# Define 3 blocks. Each block has its own normalization layer.
|
||||
# 1. Self-Attn
|
||||
if self.use_ada_layer_norm:
|
||||
self.norm1 = AdaLayerNorm(dim, num_embeds_ada_norm)
|
||||
elif self.use_ada_layer_norm_zero:
|
||||
self.norm1 = AdaLayerNormZero(dim, num_embeds_ada_norm)
|
||||
else:
|
||||
self.norm1 = nn.LayerNorm(
|
||||
dim, elementwise_affine=norm_elementwise_affine, eps=norm_eps
|
||||
)
|
||||
|
||||
self.attn1 = Attention(
|
||||
query_dim=dim,
|
||||
heads=num_attention_heads,
|
||||
dim_head=attention_head_dim,
|
||||
dropout=dropout,
|
||||
bias=attention_bias,
|
||||
upcast_attention=upcast_attention,
|
||||
)
|
||||
|
||||
# 3. Feed-forward
|
||||
if not self.use_ada_layer_norm_single:
|
||||
self.norm3 = nn.LayerNorm(
|
||||
dim, elementwise_affine=norm_elementwise_affine, eps=norm_eps
|
||||
)
|
||||
|
||||
self.ff = FeedForward(
|
||||
dim,
|
||||
dropout=dropout,
|
||||
activation_fn=activation_fn,
|
||||
final_dropout=final_dropout,
|
||||
)
|
||||
|
||||
# 4. Fuser
|
||||
if attention_type == "gated" or attention_type == "gated-text-image":
|
||||
self.fuser = GatedSelfAttentionDense(
|
||||
dim, cross_attention_dim, num_attention_heads, attention_head_dim
|
||||
)
|
||||
|
||||
# 5. Scale-shift for PixArt-Alpha.
|
||||
if self.use_ada_layer_norm_single:
|
||||
self.scale_shift_table = nn.Parameter(torch.randn(6, dim) / dim**0.5)
|
||||
|
||||
# let chunk size default to None
|
||||
self._chunk_size = None
|
||||
self._chunk_dim = 0
|
||||
|
||||
def set_chunk_feed_forward(self, chunk_size: Optional[int], dim: int = 0):
|
||||
# Sets chunk feed-forward
|
||||
self._chunk_size = chunk_size
|
||||
self._chunk_dim = dim
|
||||
|
||||
def forward(
|
||||
self,
|
||||
hidden_states: torch.FloatTensor,
|
||||
attention_mask: Optional[torch.FloatTensor] = None,
|
||||
encoder_hidden_states: Optional[torch.FloatTensor] = None,
|
||||
encoder_attention_mask: Optional[torch.FloatTensor] = None,
|
||||
timestep: Optional[torch.LongTensor] = None,
|
||||
cross_attention_kwargs: Dict[str, Any] = None,
|
||||
class_labels: Optional[torch.LongTensor] = None,
|
||||
) -> torch.FloatTensor:
|
||||
# Notice that normalization is always applied before the real computation in the following blocks.
|
||||
# 0. Self-Attention
|
||||
batch_size = hidden_states.shape[0]
|
||||
|
||||
if self.use_ada_layer_norm:
|
||||
norm_hidden_states = self.norm1(hidden_states, timestep)
|
||||
elif self.use_ada_layer_norm_zero:
|
||||
norm_hidden_states, gate_msa, shift_mlp, scale_mlp, gate_mlp = self.norm1(
|
||||
hidden_states, timestep, class_labels, hidden_dtype=hidden_states.dtype
|
||||
)
|
||||
elif self.use_layer_norm:
|
||||
norm_hidden_states = self.norm1(hidden_states)
|
||||
elif self.use_ada_layer_norm_single:
|
||||
shift_msa, scale_msa, gate_msa, shift_mlp, scale_mlp, gate_mlp = (
|
||||
self.scale_shift_table[None] + timestep.reshape(batch_size, 6, -1)
|
||||
).chunk(6, dim=1)
|
||||
norm_hidden_states = self.norm1(hidden_states)
|
||||
norm_hidden_states = norm_hidden_states * (1 + scale_msa) + shift_msa
|
||||
norm_hidden_states = norm_hidden_states.squeeze(1)
|
||||
else:
|
||||
raise ValueError("Incorrect norm used")
|
||||
|
||||
if self.pos_embed is not None:
|
||||
norm_hidden_states = self.pos_embed(norm_hidden_states)
|
||||
|
||||
# 1. Retrieve lora scale.
|
||||
lora_scale = (
|
||||
cross_attention_kwargs.get("scale", 1.0)
|
||||
if cross_attention_kwargs is not None
|
||||
else 1.0
|
||||
)
|
||||
|
||||
# 2. Prepare GLIGEN inputs
|
||||
cross_attention_kwargs = (
|
||||
cross_attention_kwargs.copy() if cross_attention_kwargs is not None else {}
|
||||
)
|
||||
gligen_kwargs = cross_attention_kwargs.pop("gligen", None)
|
||||
|
||||
attn_output = self.attn1(
|
||||
norm_hidden_states,
|
||||
attention_mask=attention_mask,
|
||||
**cross_attention_kwargs,
|
||||
)
|
||||
if self.use_ada_layer_norm_zero:
|
||||
attn_output = gate_msa.unsqueeze(1) * attn_output
|
||||
elif self.use_ada_layer_norm_single:
|
||||
attn_output = gate_msa * attn_output
|
||||
|
||||
hidden_states = attn_output + hidden_states
|
||||
if hidden_states.ndim == 4:
|
||||
hidden_states = hidden_states.squeeze(1)
|
||||
|
||||
# 2.5 GLIGEN Control
|
||||
if gligen_kwargs is not None:
|
||||
hidden_states = self.fuser(hidden_states, gligen_kwargs["objs"])
|
||||
|
||||
|
||||
# 4. Feed-forward
|
||||
if not self.use_ada_layer_norm_single:
|
||||
norm_hidden_states = self.norm3(hidden_states)
|
||||
|
||||
if self.use_ada_layer_norm_zero:
|
||||
norm_hidden_states = (
|
||||
norm_hidden_states * (1 + scale_mlp[:, None]) + shift_mlp[:, None]
|
||||
)
|
||||
|
||||
if self.use_ada_layer_norm_single:
|
||||
norm_hidden_states = self.norm2(hidden_states)
|
||||
norm_hidden_states = norm_hidden_states * (1 + scale_mlp) + shift_mlp
|
||||
|
||||
ff_output = self.ff(norm_hidden_states, scale=lora_scale)
|
||||
|
||||
if self.use_ada_layer_norm_zero:
|
||||
ff_output = gate_mlp.unsqueeze(1) * ff_output
|
||||
elif self.use_ada_layer_norm_single:
|
||||
ff_output = gate_mlp * ff_output
|
||||
|
||||
hidden_states = ff_output + hidden_states
|
||||
if hidden_states.ndim == 4:
|
||||
hidden_states = hidden_states.squeeze(1)
|
||||
|
||||
return hidden_states
|
||||
|
||||
|
||||
class TemporalBasicTransformerBlock(nn.Module):
|
||||
def __init__(
|
||||
self,
|
||||
dim: int,
|
||||
num_attention_heads: int,
|
||||
attention_head_dim: int,
|
||||
dropout=0.0,
|
||||
cross_attention_dim: Optional[int] = None,
|
||||
activation_fn: str = "geglu",
|
||||
num_embeds_ada_norm: Optional[int] = None,
|
||||
attention_bias: bool = False,
|
||||
only_cross_attention: bool = False,
|
||||
upcast_attention: bool = False,
|
||||
unet_use_cross_frame_attention=None,
|
||||
unet_use_temporal_attention=None,
|
||||
):
|
||||
super().__init__()
|
||||
self.only_cross_attention = only_cross_attention
|
||||
self.use_ada_layer_norm = num_embeds_ada_norm is not None
|
||||
self.unet_use_cross_frame_attention = unet_use_cross_frame_attention
|
||||
self.unet_use_temporal_attention = unet_use_temporal_attention
|
||||
|
||||
# SC-Attn
|
||||
self.attn1 = Attention(
|
||||
query_dim=dim,
|
||||
heads=num_attention_heads,
|
||||
dim_head=attention_head_dim,
|
||||
dropout=dropout,
|
||||
bias=attention_bias,
|
||||
upcast_attention=upcast_attention,
|
||||
)
|
||||
self.norm1 = (
|
||||
AdaLayerNorm(dim, num_embeds_ada_norm)
|
||||
if self.use_ada_layer_norm
|
||||
else nn.LayerNorm(dim)
|
||||
)
|
||||
|
||||
# Audio Cross-Attn
|
||||
if cross_attention_dim is not None:
|
||||
self.attn2 = Attention(
|
||||
query_dim=dim,
|
||||
cross_attention_dim=cross_attention_dim,
|
||||
heads=num_attention_heads,
|
||||
dim_head=attention_head_dim,
|
||||
dropout=dropout,
|
||||
bias=attention_bias,
|
||||
upcast_attention=upcast_attention,
|
||||
)
|
||||
else:
|
||||
self.attn2 = None
|
||||
|
||||
if cross_attention_dim is not None:
|
||||
self.norm2 = (
|
||||
AdaLayerNorm(dim, num_embeds_ada_norm)
|
||||
if self.use_ada_layer_norm
|
||||
else nn.LayerNorm(dim)
|
||||
)
|
||||
else:
|
||||
self.norm2 = None
|
||||
|
||||
# Feed-forward
|
||||
self.ff = FeedForward(dim, dropout=dropout, activation_fn=activation_fn)
|
||||
self.norm3 = nn.LayerNorm(dim)
|
||||
self.use_ada_layer_norm_zero = False
|
||||
|
||||
# Temp-Attn
|
||||
assert unet_use_temporal_attention is not None
|
||||
if unet_use_temporal_attention:
|
||||
self.attn_temp = Attention(
|
||||
query_dim=dim,
|
||||
heads=num_attention_heads,
|
||||
dim_head=attention_head_dim,
|
||||
dropout=dropout,
|
||||
bias=attention_bias,
|
||||
upcast_attention=upcast_attention,
|
||||
)
|
||||
nn.init.zeros_(self.attn_temp.to_out[0].weight.data)
|
||||
self.norm_temp = (
|
||||
AdaLayerNorm(dim, num_embeds_ada_norm)
|
||||
if self.use_ada_layer_norm
|
||||
else nn.LayerNorm(dim)
|
||||
)
|
||||
|
||||
def forward(
|
||||
self,
|
||||
hidden_states,
|
||||
encoder_hidden_states=None,
|
||||
audio_cond_fea = None,
|
||||
timestep=None,
|
||||
attention_mask=None,
|
||||
video_length=None,
|
||||
):
|
||||
## implemented in mutual_self_attention.py
|
||||
pass
|
||||
return hidden_states
|
||||
@@ -0,0 +1,57 @@
|
||||
from typing import Tuple
|
||||
|
||||
import torch.nn as nn
|
||||
import torch.nn.functional as F
|
||||
import torch.nn.init as init
|
||||
from diffusers.models.modeling_utils import ModelMixin
|
||||
|
||||
from .motion_module import zero_module
|
||||
from .resnet import InflatedConv3d
|
||||
|
||||
|
||||
class FaceLocator(ModelMixin):
|
||||
def __init__(
|
||||
self,
|
||||
conditioning_embedding_channels: int,
|
||||
conditioning_channels: int = 1,
|
||||
block_out_channels: Tuple[int] = (16, 32, 64, 128),
|
||||
):
|
||||
super().__init__()
|
||||
self.conv_in = InflatedConv3d(
|
||||
conditioning_channels, block_out_channels[0], kernel_size=3, padding=1
|
||||
)
|
||||
|
||||
self.blocks = nn.ModuleList([])
|
||||
|
||||
for i in range(len(block_out_channels) - 1):
|
||||
channel_in = block_out_channels[i]
|
||||
channel_out = block_out_channels[i + 1]
|
||||
self.blocks.append(
|
||||
InflatedConv3d(channel_in, channel_in, kernel_size=3, padding=1)
|
||||
)
|
||||
self.blocks.append(
|
||||
InflatedConv3d(
|
||||
channel_in, channel_out, kernel_size=3, padding=1, stride=2
|
||||
)
|
||||
)
|
||||
|
||||
self.conv_out = zero_module(
|
||||
InflatedConv3d(
|
||||
block_out_channels[-1],
|
||||
conditioning_embedding_channels,
|
||||
kernel_size=3,
|
||||
padding=1,
|
||||
)
|
||||
)
|
||||
|
||||
def forward(self, conditioning):
|
||||
embedding = self.conv_in(conditioning)
|
||||
embedding = F.silu(embedding)
|
||||
|
||||
for block in self.blocks:
|
||||
embedding = block(embedding)
|
||||
embedding = F.silu(embedding)
|
||||
|
||||
embedding = self.conv_out(embedding)
|
||||
|
||||
return embedding
|
||||
@@ -0,0 +1,481 @@
|
||||
# Adapt from https://github.com/guoyww/AnimateDiff/blob/main/animatediff/models/motion_module.py
|
||||
import math
|
||||
from dataclasses import dataclass
|
||||
from pathlib import Path
|
||||
from typing import Callable, Optional
|
||||
|
||||
import torch
|
||||
from diffusers.models.attention import FeedForward
|
||||
from diffusers.models.attention_processor import Attention, AttnProcessor
|
||||
from diffusers.utils import BaseOutput
|
||||
from diffusers.utils.import_utils import is_xformers_available
|
||||
from einops import rearrange, repeat
|
||||
from torch import nn
|
||||
import cv2
|
||||
import numpy as np
|
||||
|
||||
def zero_module(module):
|
||||
# Zero out the parameters of a module and return it.
|
||||
for p in module.parameters():
|
||||
p.detach().zero_()
|
||||
return module
|
||||
|
||||
|
||||
@dataclass
|
||||
class TemporalTransformer3DModelOutput(BaseOutput):
|
||||
sample: torch.FloatTensor
|
||||
|
||||
|
||||
if is_xformers_available():
|
||||
import xformers
|
||||
import xformers.ops
|
||||
else:
|
||||
xformers = None
|
||||
|
||||
|
||||
def get_motion_module(in_channels, motion_module_type: str, motion_module_kwargs: dict):
|
||||
if motion_module_type == "Vanilla":
|
||||
return VanillaTemporalModule(
|
||||
in_channels=in_channels,
|
||||
**motion_module_kwargs,
|
||||
)
|
||||
else:
|
||||
raise ValueError
|
||||
|
||||
|
||||
class VanillaTemporalModule(nn.Module):
|
||||
def __init__(
|
||||
self,
|
||||
in_channels,
|
||||
num_attention_heads=8,
|
||||
num_transformer_block=2,
|
||||
attention_block_types=("Temporal_Self", "Temporal_Self"),
|
||||
cross_frame_attention_mode=None,
|
||||
temporal_position_encoding=False,
|
||||
temporal_position_encoding_max_len=24,
|
||||
temporal_attention_dim_div=1,
|
||||
zero_initialize=True,
|
||||
):
|
||||
super().__init__()
|
||||
|
||||
b = 1
|
||||
f = 12
|
||||
c_64, c_32, c_16, c_8 = 320 ,640 ,1280, 1280
|
||||
masks_dir = Path(__file__).resolve().parents[2] / "masks"
|
||||
self.mask_64 = self.create_mask(masks_dir / "mask_64.png", b, c_64, f)
|
||||
self.mask_32 = self.create_mask(masks_dir / "mask_32.png", b, c_32, f)
|
||||
self.mask_16 = self.create_mask(masks_dir / "mask_16.png", b, c_16, f)
|
||||
self.mask_8 = self.create_mask(masks_dir / "mask_8.png", b, c_8, f)
|
||||
|
||||
self.temporal_transformer = TemporalTransformer3DModel(
|
||||
in_channels=in_channels,
|
||||
num_attention_heads=num_attention_heads,
|
||||
attention_head_dim=in_channels
|
||||
// num_attention_heads
|
||||
// temporal_attention_dim_div,
|
||||
num_layers=num_transformer_block,
|
||||
attention_block_types=attention_block_types,
|
||||
cross_frame_attention_mode=cross_frame_attention_mode,
|
||||
temporal_position_encoding=temporal_position_encoding,
|
||||
temporal_position_encoding_max_len=temporal_position_encoding_max_len,
|
||||
)
|
||||
|
||||
if zero_initialize:
|
||||
self.temporal_transformer.proj_out = zero_module(
|
||||
self.temporal_transformer.proj_out
|
||||
)
|
||||
|
||||
def create_mask(self, mask_path, b, c_new, f):
|
||||
raw_mask = cv2.imread(str(mask_path))
|
||||
if raw_mask is None:
|
||||
raise FileNotFoundError(f"HighSync motion mask could not be read: {mask_path}")
|
||||
mask = torch.FloatTensor(1 - (raw_mask / 255.))
|
||||
(h, w, c) = mask.shape
|
||||
mask_new = torch.zeros((h,w,c_new))
|
||||
mask_new[:, :, :c] = mask
|
||||
mask_new[:, :, c:] = mask[:, :, c-1:c].repeat(1,1,c_new-c)
|
||||
|
||||
mask_new = mask_new.permute(2,0,1).unsqueeze(1).unsqueeze(0) # 1*c*1*h*w
|
||||
|
||||
mask_new = mask_new.repeat(b, 1, f, 1, 1) # b*c*f*h*w
|
||||
|
||||
device = "cuda" if torch.cuda.is_available() else "cpu"
|
||||
return mask_new.to(device, dtype=torch.float16)
|
||||
|
||||
|
||||
def forward(
|
||||
self,
|
||||
input_tensor,
|
||||
temb,
|
||||
encoder_hidden_states,
|
||||
attention_mask=None,
|
||||
anchor_frame_idx=None,
|
||||
):
|
||||
|
||||
b, c, f, h, w = input_tensor.shape
|
||||
|
||||
if h==64:
|
||||
mask = self.mask_64
|
||||
elif h==32:
|
||||
mask = self.mask_32
|
||||
elif h==16:
|
||||
mask = self.mask_16
|
||||
elif h==8:
|
||||
mask = self.mask_8
|
||||
|
||||
hidden_states = input_tensor * mask
|
||||
hidden_states = self.temporal_transformer(
|
||||
hidden_states, encoder_hidden_states, attention_mask
|
||||
)
|
||||
|
||||
output = hidden_states*mask + input_tensor*(1-mask)
|
||||
# output = input_tensor
|
||||
|
||||
return output
|
||||
|
||||
|
||||
# class VanillaTemporalModule(nn.Module):
|
||||
# def __init__(
|
||||
# self,
|
||||
# in_channels,
|
||||
# num_attention_heads=8,
|
||||
# num_transformer_block=2,
|
||||
# attention_block_types=("Temporal_Self", "Temporal_Self"),
|
||||
# cross_frame_attention_mode=None,
|
||||
# temporal_position_encoding=False,
|
||||
# temporal_position_encoding_max_len=24,
|
||||
# temporal_attention_dim_div=1,
|
||||
# zero_initialize=True,
|
||||
# ):
|
||||
# super().__init__()
|
||||
|
||||
# self.temporal_transformer = TemporalTransformer3DModel(
|
||||
# in_channels=in_channels,
|
||||
# num_attention_heads=num_attention_heads,
|
||||
# attention_head_dim=in_channels
|
||||
# // num_attention_heads
|
||||
# // temporal_attention_dim_div,
|
||||
# num_layers=num_transformer_block,
|
||||
# attention_block_types=attention_block_types,
|
||||
# cross_frame_attention_mode=cross_frame_attention_mode,
|
||||
# temporal_position_encoding=temporal_position_encoding,
|
||||
# temporal_position_encoding_max_len=temporal_position_encoding_max_len,
|
||||
# )
|
||||
|
||||
# if zero_initialize:
|
||||
# self.temporal_transformer.proj_out = zero_module(
|
||||
# self.temporal_transformer.proj_out
|
||||
# )
|
||||
|
||||
# def forward(
|
||||
# self,
|
||||
# input_tensor,
|
||||
# temb,
|
||||
# encoder_hidden_states,
|
||||
# attention_mask=None,
|
||||
# anchor_frame_idx=None,
|
||||
# ):
|
||||
# hidden_states = input_tensor
|
||||
# hidden_states = self.temporal_transformer(
|
||||
# hidden_states, encoder_hidden_states, attention_mask
|
||||
# )
|
||||
|
||||
# output = hidden_states
|
||||
# return output
|
||||
|
||||
|
||||
class TemporalTransformer3DModel(nn.Module):
|
||||
def __init__(
|
||||
self,
|
||||
in_channels,
|
||||
num_attention_heads,
|
||||
attention_head_dim,
|
||||
num_layers,
|
||||
attention_block_types=(
|
||||
"Temporal_Self",
|
||||
"Temporal_Self",
|
||||
),
|
||||
dropout=0.0,
|
||||
norm_num_groups=32,
|
||||
cross_attention_dim=768,
|
||||
activation_fn="geglu",
|
||||
attention_bias=False,
|
||||
upcast_attention=False,
|
||||
cross_frame_attention_mode=None,
|
||||
temporal_position_encoding=False,
|
||||
temporal_position_encoding_max_len=24,
|
||||
):
|
||||
super().__init__()
|
||||
|
||||
inner_dim = num_attention_heads * attention_head_dim
|
||||
|
||||
self.norm = torch.nn.GroupNorm(
|
||||
num_groups=norm_num_groups, num_channels=in_channels, eps=1e-6, affine=True
|
||||
)
|
||||
self.proj_in = nn.Linear(in_channels, inner_dim)
|
||||
|
||||
self.transformer_blocks = nn.ModuleList(
|
||||
[
|
||||
TemporalTransformerBlock(
|
||||
dim=inner_dim,
|
||||
num_attention_heads=num_attention_heads,
|
||||
attention_head_dim=attention_head_dim,
|
||||
attention_block_types=attention_block_types,
|
||||
dropout=dropout,
|
||||
norm_num_groups=norm_num_groups,
|
||||
cross_attention_dim=cross_attention_dim,
|
||||
activation_fn=activation_fn,
|
||||
attention_bias=attention_bias,
|
||||
upcast_attention=upcast_attention,
|
||||
cross_frame_attention_mode=cross_frame_attention_mode,
|
||||
temporal_position_encoding=temporal_position_encoding,
|
||||
temporal_position_encoding_max_len=temporal_position_encoding_max_len,
|
||||
)
|
||||
for d in range(num_layers)
|
||||
]
|
||||
)
|
||||
self.proj_out = nn.Linear(inner_dim, in_channels)
|
||||
|
||||
def forward(self, hidden_states, encoder_hidden_states=None, attention_mask=None):
|
||||
assert (
|
||||
hidden_states.dim() == 5
|
||||
), f"Expected hidden_states to have ndim=5, but got ndim={hidden_states.dim()}."
|
||||
video_length = hidden_states.shape[2]
|
||||
hidden_states = rearrange(hidden_states, "b c f h w -> (b f) c h w")
|
||||
|
||||
batch, channel, height, weight = hidden_states.shape
|
||||
residual = hidden_states
|
||||
|
||||
hidden_states = self.norm(hidden_states)
|
||||
inner_dim = hidden_states.shape[1]
|
||||
hidden_states = hidden_states.permute(0, 2, 3, 1).reshape(
|
||||
batch, height * weight, inner_dim
|
||||
)
|
||||
hidden_states = self.proj_in(hidden_states)
|
||||
|
||||
# Transformer Blocks
|
||||
for block in self.transformer_blocks:
|
||||
hidden_states = block(
|
||||
hidden_states,
|
||||
encoder_hidden_states=encoder_hidden_states,
|
||||
video_length=video_length,
|
||||
)
|
||||
|
||||
# output
|
||||
hidden_states = self.proj_out(hidden_states)
|
||||
hidden_states = (
|
||||
hidden_states.reshape(batch, height, weight, inner_dim)
|
||||
.permute(0, 3, 1, 2)
|
||||
.contiguous()
|
||||
)
|
||||
|
||||
output = hidden_states + residual
|
||||
output = rearrange(output, "(b f) c h w -> b c f h w", f=video_length)
|
||||
|
||||
return output
|
||||
|
||||
|
||||
class TemporalTransformerBlock(nn.Module):
|
||||
def __init__(
|
||||
self,
|
||||
dim,
|
||||
num_attention_heads,
|
||||
attention_head_dim,
|
||||
attention_block_types=(
|
||||
"Temporal_Self",
|
||||
"Temporal_Self",
|
||||
),
|
||||
dropout=0.0,
|
||||
norm_num_groups=32,
|
||||
cross_attention_dim=768,
|
||||
activation_fn="geglu",
|
||||
attention_bias=False,
|
||||
upcast_attention=False,
|
||||
cross_frame_attention_mode=None,
|
||||
temporal_position_encoding=False,
|
||||
temporal_position_encoding_max_len=24,
|
||||
):
|
||||
super().__init__()
|
||||
|
||||
attention_blocks = []
|
||||
norms = []
|
||||
|
||||
for block_name in attention_block_types:
|
||||
attention_blocks.append(
|
||||
VersatileAttention(
|
||||
attention_mode=block_name.split("_")[0],
|
||||
cross_attention_dim=cross_attention_dim
|
||||
if block_name.endswith("_Cross")
|
||||
else None,
|
||||
query_dim=dim,
|
||||
heads=num_attention_heads,
|
||||
dim_head=attention_head_dim,
|
||||
dropout=dropout,
|
||||
bias=attention_bias,
|
||||
upcast_attention=upcast_attention,
|
||||
cross_frame_attention_mode=cross_frame_attention_mode,
|
||||
temporal_position_encoding=temporal_position_encoding,
|
||||
temporal_position_encoding_max_len=temporal_position_encoding_max_len,
|
||||
)
|
||||
)
|
||||
norms.append(nn.LayerNorm(dim))
|
||||
|
||||
self.attention_blocks = nn.ModuleList(attention_blocks)
|
||||
self.norms = nn.ModuleList(norms)
|
||||
|
||||
self.ff = FeedForward(dim, dropout=dropout, activation_fn=activation_fn)
|
||||
self.ff_norm = nn.LayerNorm(dim)
|
||||
|
||||
def forward(
|
||||
self,
|
||||
hidden_states,
|
||||
encoder_hidden_states=None,
|
||||
attention_mask=None,
|
||||
video_length=None,
|
||||
):
|
||||
for attention_block, norm in zip(self.attention_blocks, self.norms):
|
||||
norm_hidden_states = norm(hidden_states)
|
||||
hidden_states = (
|
||||
attention_block(
|
||||
norm_hidden_states,
|
||||
encoder_hidden_states=encoder_hidden_states
|
||||
if attention_block.is_cross_attention
|
||||
else None,
|
||||
video_length=video_length,
|
||||
)
|
||||
+ hidden_states
|
||||
)
|
||||
|
||||
hidden_states = self.ff(self.ff_norm(hidden_states)) + hidden_states
|
||||
|
||||
output = hidden_states
|
||||
return output
|
||||
|
||||
|
||||
class PositionalEncoding(nn.Module):
|
||||
def __init__(self, d_model, dropout=0.0, max_len=24):
|
||||
super().__init__()
|
||||
self.dropout = nn.Dropout(p=dropout)
|
||||
position = torch.arange(max_len).unsqueeze(1)
|
||||
div_term = torch.exp(
|
||||
torch.arange(0, d_model, 2) * (-math.log(10000.0) / d_model)
|
||||
)
|
||||
pe = torch.zeros(1, max_len, d_model)
|
||||
pe[0, :, 0::2] = torch.sin(position * div_term)
|
||||
pe[0, :, 1::2] = torch.cos(position * div_term)
|
||||
self.register_buffer("pe", pe)
|
||||
|
||||
def forward(self, x):
|
||||
x = x + self.pe[:, : x.size(1)]
|
||||
return self.dropout(x)
|
||||
|
||||
|
||||
class VersatileAttention(Attention):
|
||||
def __init__(
|
||||
self,
|
||||
attention_mode=None,
|
||||
cross_frame_attention_mode=None,
|
||||
temporal_position_encoding=False,
|
||||
temporal_position_encoding_max_len=24,
|
||||
*args,
|
||||
**kwargs,
|
||||
):
|
||||
super().__init__(*args, **kwargs)
|
||||
assert attention_mode == "Temporal"
|
||||
|
||||
self.attention_mode = attention_mode
|
||||
self.is_cross_attention = kwargs["cross_attention_dim"] is not None
|
||||
|
||||
self.pos_encoder = (
|
||||
PositionalEncoding(
|
||||
kwargs["query_dim"],
|
||||
dropout=0.0,
|
||||
max_len=temporal_position_encoding_max_len,
|
||||
)
|
||||
if (temporal_position_encoding and attention_mode == "Temporal")
|
||||
else None
|
||||
)
|
||||
|
||||
def extra_repr(self):
|
||||
return f"(Module Info) Attention_Mode: {self.attention_mode}, Is_Cross_Attention: {self.is_cross_attention}"
|
||||
|
||||
def set_use_memory_efficient_attention_xformers(
|
||||
self,
|
||||
use_memory_efficient_attention_xformers: bool,
|
||||
attention_op: Optional[Callable] = None,
|
||||
):
|
||||
if use_memory_efficient_attention_xformers:
|
||||
if not is_xformers_available():
|
||||
raise ModuleNotFoundError(
|
||||
(
|
||||
"Refer to https://github.com/facebookresearch/xformers for more information on how to install"
|
||||
" xformers"
|
||||
),
|
||||
name="xformers",
|
||||
)
|
||||
elif not torch.cuda.is_available():
|
||||
raise ValueError(
|
||||
"torch.cuda.is_available() should be True but is False. xformers' memory efficient attention is"
|
||||
" only available for GPU "
|
||||
)
|
||||
else:
|
||||
try:
|
||||
# Make sure we can run the memory efficient attention
|
||||
_ = xformers.ops.memory_efficient_attention(
|
||||
torch.randn((1, 2, 40), device="cuda"),
|
||||
torch.randn((1, 2, 40), device="cuda"),
|
||||
torch.randn((1, 2, 40), device="cuda"),
|
||||
)
|
||||
except Exception as e:
|
||||
raise e
|
||||
|
||||
# XFormersAttnProcessor corrupts video generation and work with Pytorch 1.13.
|
||||
# Pytorch 2.0.1 AttnProcessor works the same as XFormersAttnProcessor in Pytorch 1.13.
|
||||
# You don't need XFormersAttnProcessor here.
|
||||
# processor = XFormersAttnProcessor(
|
||||
# attention_op=attention_op,
|
||||
# )
|
||||
processor = AttnProcessor()
|
||||
else:
|
||||
processor = AttnProcessor()
|
||||
|
||||
self.set_processor(processor)
|
||||
|
||||
def forward(
|
||||
self,
|
||||
hidden_states,
|
||||
encoder_hidden_states=None,
|
||||
attention_mask=None,
|
||||
video_length=None,
|
||||
**cross_attention_kwargs,
|
||||
):
|
||||
if self.attention_mode == "Temporal":
|
||||
d = hidden_states.shape[1] # d means HxW
|
||||
hidden_states = rearrange(
|
||||
hidden_states, "(b f) d c -> (b d) f c", f=video_length
|
||||
)
|
||||
|
||||
if self.pos_encoder is not None:
|
||||
hidden_states = self.pos_encoder(hidden_states)
|
||||
|
||||
encoder_hidden_states = (
|
||||
repeat(encoder_hidden_states, "b n c -> (b d) n c", d=d)
|
||||
if encoder_hidden_states is not None
|
||||
else encoder_hidden_states
|
||||
)
|
||||
|
||||
else:
|
||||
raise NotImplementedError
|
||||
|
||||
hidden_states = self.processor(
|
||||
self,
|
||||
hidden_states,
|
||||
encoder_hidden_states=encoder_hidden_states,
|
||||
attention_mask=attention_mask,
|
||||
**cross_attention_kwargs,
|
||||
)
|
||||
|
||||
if self.attention_mode == "Temporal":
|
||||
hidden_states = rearrange(hidden_states, "(b d) f c -> (b f) d c", d=d)
|
||||
|
||||
return hidden_states
|
||||
@@ -0,0 +1,336 @@
|
||||
# Adapted from https://github.com/magic-research/magic-animate/blob/main/magicanimate/models/mutual_self_attention.py
|
||||
from typing import Any, Dict, Optional
|
||||
|
||||
import torch
|
||||
from einops import rearrange
|
||||
|
||||
from .attention import TemporalBasicTransformerBlock
|
||||
|
||||
from .attention import BasicTransformerBlock
|
||||
|
||||
|
||||
def torch_dfs(model: torch.nn.Module):
|
||||
result = [model]
|
||||
for child in model.children():
|
||||
result += torch_dfs(child)
|
||||
return result
|
||||
|
||||
|
||||
class ReferenceAttentionControl:
|
||||
def __init__(
|
||||
self,
|
||||
unet,
|
||||
mode="write",
|
||||
do_classifier_free_guidance=False,
|
||||
attention_auto_machine_weight=float("inf"),
|
||||
gn_auto_machine_weight=1.0,
|
||||
style_fidelity=1.0,
|
||||
reference_attn=True,
|
||||
reference_adain=False,
|
||||
fusion_blocks="midup",
|
||||
batch_size=1,
|
||||
audio_feature_ratio = 3.0
|
||||
) -> None:
|
||||
# 10. Modify self attention and group norm
|
||||
self.unet = unet
|
||||
self.audio_feature_ratio = audio_feature_ratio
|
||||
assert mode in ["read", "write"]
|
||||
assert fusion_blocks in ["midup", "full"]
|
||||
self.reference_attn = reference_attn
|
||||
self.reference_adain = reference_adain
|
||||
self.fusion_blocks = fusion_blocks
|
||||
self.register_reference_hooks(
|
||||
mode,
|
||||
do_classifier_free_guidance,
|
||||
attention_auto_machine_weight,
|
||||
gn_auto_machine_weight,
|
||||
style_fidelity,
|
||||
reference_attn,
|
||||
reference_adain,
|
||||
fusion_blocks,
|
||||
batch_size=batch_size,
|
||||
)
|
||||
|
||||
def register_reference_hooks(
|
||||
self,
|
||||
mode,
|
||||
do_classifier_free_guidance,
|
||||
attention_auto_machine_weight,
|
||||
gn_auto_machine_weight,
|
||||
style_fidelity,
|
||||
reference_attn,
|
||||
reference_adain,
|
||||
dtype=torch.float16,
|
||||
batch_size=1,
|
||||
num_images_per_prompt=1,
|
||||
device=torch.device("cuda"),
|
||||
fusion_blocks="midup",
|
||||
):
|
||||
MODE = mode
|
||||
audio_feature_ratio = self.audio_feature_ratio
|
||||
do_classifier_free_guidance = do_classifier_free_guidance
|
||||
attention_auto_machine_weight = attention_auto_machine_weight
|
||||
gn_auto_machine_weight = gn_auto_machine_weight
|
||||
style_fidelity = style_fidelity
|
||||
reference_attn = reference_attn
|
||||
reference_adain = reference_adain
|
||||
fusion_blocks = fusion_blocks
|
||||
num_images_per_prompt = num_images_per_prompt
|
||||
dtype = dtype
|
||||
if do_classifier_free_guidance:
|
||||
uc_mask = (
|
||||
torch.Tensor(
|
||||
[1] * batch_size * num_images_per_prompt * 16
|
||||
+ [0] * batch_size * num_images_per_prompt * 16
|
||||
)
|
||||
.to(device)
|
||||
.bool()
|
||||
)
|
||||
else:
|
||||
uc_mask = (
|
||||
torch.Tensor([0] * batch_size * num_images_per_prompt * 2)
|
||||
.to(device)
|
||||
.bool()
|
||||
)
|
||||
|
||||
def hacked_basic_transformer_inner_forward(
|
||||
self,
|
||||
hidden_states: torch.FloatTensor,
|
||||
attention_mask: Optional[torch.FloatTensor] = None,
|
||||
audio_cond_fea: Optional[torch.FloatTensor] = None,
|
||||
encoder_hidden_states: Optional[torch.FloatTensor] = None,
|
||||
encoder_attention_mask: Optional[torch.FloatTensor] = None,
|
||||
timestep: Optional[torch.LongTensor] = None,
|
||||
cross_attention_kwargs: Dict[str, Any] = None,
|
||||
class_labels: Optional[torch.LongTensor] = None,
|
||||
video_length=None,
|
||||
):
|
||||
if self.use_ada_layer_norm: # False
|
||||
norm_hidden_states = self.norm1(hidden_states, timestep)
|
||||
elif self.use_ada_layer_norm_zero:
|
||||
(
|
||||
norm_hidden_states,
|
||||
gate_msa,
|
||||
shift_mlp,
|
||||
scale_mlp,
|
||||
gate_mlp,
|
||||
) = self.norm1(
|
||||
hidden_states,
|
||||
timestep,
|
||||
class_labels,
|
||||
hidden_dtype=hidden_states.dtype,
|
||||
)
|
||||
else:
|
||||
norm_hidden_states = self.norm1(hidden_states)
|
||||
|
||||
# 1. Self-Attention
|
||||
# self.only_cross_attention = False
|
||||
cross_attention_kwargs = (
|
||||
cross_attention_kwargs if cross_attention_kwargs is not None else {}
|
||||
)
|
||||
if self.only_cross_attention:
|
||||
attn_output = self.attn1(
|
||||
norm_hidden_states,
|
||||
attention_mask=attention_mask,
|
||||
**cross_attention_kwargs,
|
||||
)
|
||||
else:
|
||||
if MODE == "write":
|
||||
self.bank.append(norm_hidden_states.clone())
|
||||
attn_output = self.attn1(
|
||||
norm_hidden_states,
|
||||
attention_mask=attention_mask,
|
||||
**cross_attention_kwargs,
|
||||
)
|
||||
if MODE == "read":
|
||||
# bank_feas = [
|
||||
# rearrange(
|
||||
# d.unsqueeze(1).repeat(1, video_length, 1, 1),
|
||||
# "b t l c -> (b t) l c",
|
||||
# )
|
||||
# for d in self.bank
|
||||
# ]
|
||||
|
||||
bank_feas = [d
|
||||
for d in self.bank
|
||||
]
|
||||
|
||||
modify_norm_hidden_states = torch.cat(
|
||||
[norm_hidden_states] + bank_feas, dim=1
|
||||
)
|
||||
# Reference Cross-Attention
|
||||
hidden_states_uc = (
|
||||
self.attn1(
|
||||
norm_hidden_states,
|
||||
encoder_hidden_states=modify_norm_hidden_states,
|
||||
attention_mask=attention_mask,
|
||||
)
|
||||
+ hidden_states
|
||||
)
|
||||
if do_classifier_free_guidance:
|
||||
## inference with cfg
|
||||
hidden_states_c = hidden_states_uc.clone()
|
||||
_uc_mask = uc_mask.clone()
|
||||
if hidden_states.shape[0] != _uc_mask.shape[0]:
|
||||
_uc_mask = (
|
||||
torch.Tensor(
|
||||
[1] * (hidden_states.shape[0] // 2)
|
||||
+ [0] * (hidden_states.shape[0] // 2)
|
||||
)
|
||||
.to(device)
|
||||
.bool()
|
||||
)
|
||||
hidden_states_c[_uc_mask] = (
|
||||
self.attn1(
|
||||
norm_hidden_states[_uc_mask],
|
||||
encoder_hidden_states=norm_hidden_states[_uc_mask], # B * 4096 * 768
|
||||
attention_mask=attention_mask,
|
||||
)
|
||||
+ hidden_states[_uc_mask]
|
||||
)
|
||||
hidden_states = hidden_states_c.clone()
|
||||
else:
|
||||
hidden_states = hidden_states_uc
|
||||
|
||||
if self.attn2 is not None:
|
||||
# Audio Cross-Attention
|
||||
norm_hidden_states = (
|
||||
self.norm2(hidden_states, timestep)
|
||||
if self.use_ada_layer_norm
|
||||
else self.norm2(hidden_states)
|
||||
)
|
||||
hidden_states = (
|
||||
self.attn2(
|
||||
norm_hidden_states,
|
||||
encoder_hidden_states=audio_cond_fea,
|
||||
attention_mask=attention_mask,
|
||||
) * audio_feature_ratio
|
||||
+ hidden_states
|
||||
)
|
||||
|
||||
# Feed-forward
|
||||
hidden_states = self.ff(self.norm3(hidden_states)) + hidden_states
|
||||
|
||||
# Temporal-Attention
|
||||
return hidden_states
|
||||
|
||||
if self.use_ada_layer_norm_zero:
|
||||
attn_output = gate_msa.unsqueeze(1) * attn_output
|
||||
hidden_states = attn_output + hidden_states
|
||||
|
||||
# 3. Feed-forward
|
||||
norm_hidden_states = self.norm3(hidden_states)
|
||||
|
||||
if self.use_ada_layer_norm_zero:
|
||||
norm_hidden_states = (
|
||||
norm_hidden_states * (1 + scale_mlp[:, None]) + shift_mlp[:, None]
|
||||
)
|
||||
|
||||
ff_output = self.ff(norm_hidden_states)
|
||||
|
||||
if self.use_ada_layer_norm_zero:
|
||||
ff_output = gate_mlp.unsqueeze(1) * ff_output
|
||||
|
||||
hidden_states = ff_output + hidden_states
|
||||
|
||||
return hidden_states
|
||||
|
||||
if self.reference_attn:
|
||||
if self.fusion_blocks == "midup":
|
||||
attn_modules = [
|
||||
module
|
||||
for module in (
|
||||
torch_dfs(self.unet.mid_block) + torch_dfs(self.unet.up_blocks)
|
||||
)
|
||||
if isinstance(module, BasicTransformerBlock)
|
||||
or isinstance(module, TemporalBasicTransformerBlock)
|
||||
]
|
||||
elif self.fusion_blocks == "full":
|
||||
attn_modules = [
|
||||
module
|
||||
for module in torch_dfs(self.unet)
|
||||
if isinstance(module, BasicTransformerBlock)
|
||||
or isinstance(module, TemporalBasicTransformerBlock)
|
||||
]
|
||||
attn_modules = sorted(
|
||||
attn_modules, key=lambda x: -x.norm1.normalized_shape[0]
|
||||
)
|
||||
|
||||
for i, module in enumerate(attn_modules):
|
||||
module._original_inner_forward = module.forward
|
||||
if isinstance(module, BasicTransformerBlock):
|
||||
module.forward = hacked_basic_transformer_inner_forward.__get__(
|
||||
module, BasicTransformerBlock
|
||||
)
|
||||
if isinstance(module, TemporalBasicTransformerBlock):
|
||||
module.forward = hacked_basic_transformer_inner_forward.__get__(
|
||||
module, TemporalBasicTransformerBlock
|
||||
)
|
||||
|
||||
module.bank = []
|
||||
module.attn_weight = float(i) / float(len(attn_modules))
|
||||
|
||||
def update(self, writer, do_classifier_free_guidance=False, dtype=torch.float16):
|
||||
if self.reference_attn:
|
||||
if self.fusion_blocks == "midup":
|
||||
reader_attn_modules = [
|
||||
module
|
||||
for module in (
|
||||
torch_dfs(self.unet.mid_block) + torch_dfs(self.unet.up_blocks)
|
||||
)
|
||||
if isinstance(module, TemporalBasicTransformerBlock)
|
||||
]
|
||||
writer_attn_modules = [
|
||||
module
|
||||
for module in (
|
||||
torch_dfs(writer.unet.mid_block)
|
||||
+ torch_dfs(writer.unet.up_blocks)
|
||||
)
|
||||
if isinstance(module, BasicTransformerBlock)
|
||||
]
|
||||
elif self.fusion_blocks == "full":
|
||||
reader_attn_modules = [
|
||||
module
|
||||
for module in torch_dfs(self.unet)
|
||||
if isinstance(module, TemporalBasicTransformerBlock)
|
||||
]
|
||||
writer_attn_modules = [
|
||||
module
|
||||
for module in torch_dfs(writer.unet)
|
||||
if isinstance(module, BasicTransformerBlock)
|
||||
]
|
||||
reader_attn_modules = sorted(
|
||||
reader_attn_modules, key=lambda x: -x.norm1.normalized_shape[0]
|
||||
)
|
||||
writer_attn_modules = sorted(
|
||||
writer_attn_modules, key=lambda x: -x.norm1.normalized_shape[0]
|
||||
)
|
||||
for r, w in zip(reader_attn_modules, writer_attn_modules):
|
||||
if do_classifier_free_guidance:
|
||||
r.bank = [torch.cat([v, v]).to(dtype) for v in w.bank]
|
||||
else:
|
||||
r.bank = [v.clone().to(dtype) for v in w.bank]
|
||||
|
||||
def clear(self):
|
||||
if self.reference_attn:
|
||||
if self.fusion_blocks == "midup":
|
||||
reader_attn_modules = [
|
||||
module
|
||||
for module in (
|
||||
torch_dfs(self.unet.mid_block) + torch_dfs(self.unet.up_blocks)
|
||||
)
|
||||
if isinstance(module, BasicTransformerBlock)
|
||||
or isinstance(module, TemporalBasicTransformerBlock)
|
||||
]
|
||||
elif self.fusion_blocks == "full":
|
||||
reader_attn_modules = [
|
||||
module
|
||||
for module in torch_dfs(self.unet)
|
||||
if isinstance(module, BasicTransformerBlock)
|
||||
or isinstance(module, TemporalBasicTransformerBlock)
|
||||
]
|
||||
reader_attn_modules = sorted(
|
||||
reader_attn_modules, key=lambda x: -x.norm1.normalized_shape[0]
|
||||
)
|
||||
for r in reader_attn_modules:
|
||||
r.bank.clear()
|
||||
@@ -0,0 +1,252 @@
|
||||
# Adapted from https://github.com/huggingface/diffusers/blob/main/src/diffusers/models/resnet.py
|
||||
|
||||
import torch
|
||||
import torch.nn as nn
|
||||
import torch.nn.functional as F
|
||||
from einops import rearrange
|
||||
|
||||
|
||||
class InflatedConv3d(nn.Conv2d):
|
||||
def forward(self, x):
|
||||
video_length = x.shape[2]
|
||||
|
||||
x = rearrange(x, "b c f h w -> (b f) c h w")
|
||||
x = super().forward(x)
|
||||
x = rearrange(x, "(b f) c h w -> b c f h w", f=video_length)
|
||||
|
||||
return x
|
||||
|
||||
|
||||
class InflatedGroupNorm(nn.GroupNorm):
|
||||
def forward(self, x):
|
||||
video_length = x.shape[2]
|
||||
|
||||
x = rearrange(x, "b c f h w -> (b f) c h w")
|
||||
x = super().forward(x)
|
||||
x = rearrange(x, "(b f) c h w -> b c f h w", f=video_length)
|
||||
|
||||
return x
|
||||
|
||||
|
||||
class Upsample3D(nn.Module):
|
||||
def __init__(
|
||||
self,
|
||||
channels,
|
||||
use_conv=False,
|
||||
use_conv_transpose=False,
|
||||
out_channels=None,
|
||||
name="conv",
|
||||
):
|
||||
super().__init__()
|
||||
self.channels = channels
|
||||
self.out_channels = out_channels or channels
|
||||
self.use_conv = use_conv
|
||||
self.use_conv_transpose = use_conv_transpose
|
||||
self.name = name
|
||||
|
||||
conv = None
|
||||
if use_conv_transpose:
|
||||
raise NotImplementedError
|
||||
elif use_conv:
|
||||
self.conv = InflatedConv3d(self.channels, self.out_channels, 3, padding=1)
|
||||
|
||||
def forward(self, hidden_states, output_size=None):
|
||||
assert hidden_states.shape[1] == self.channels
|
||||
|
||||
if self.use_conv_transpose:
|
||||
raise NotImplementedError
|
||||
|
||||
# Cast to float32 to as 'upsample_nearest2d_out_frame' op does not support bfloat16
|
||||
dtype = hidden_states.dtype
|
||||
if dtype == torch.bfloat16:
|
||||
hidden_states = hidden_states.to(torch.float32)
|
||||
|
||||
# upsample_nearest_nhwc fails with large batch sizes. see https://github.com/huggingface/diffusers/issues/984
|
||||
if hidden_states.shape[0] >= 64:
|
||||
hidden_states = hidden_states.contiguous()
|
||||
|
||||
# if `output_size` is passed we force the interpolation output
|
||||
# size and do not make use of `scale_factor=2`
|
||||
if output_size is None:
|
||||
hidden_states = F.interpolate(
|
||||
hidden_states, scale_factor=[1.0, 2.0, 2.0], mode="nearest"
|
||||
)
|
||||
else:
|
||||
hidden_states = F.interpolate(
|
||||
hidden_states, size=output_size, mode="nearest"
|
||||
)
|
||||
|
||||
# If the input is bfloat16, we cast back to bfloat16
|
||||
if dtype == torch.bfloat16:
|
||||
hidden_states = hidden_states.to(dtype)
|
||||
|
||||
# if self.use_conv:
|
||||
# if self.name == "conv":
|
||||
# hidden_states = self.conv(hidden_states)
|
||||
# else:
|
||||
# hidden_states = self.Conv2d_0(hidden_states)
|
||||
hidden_states = self.conv(hidden_states)
|
||||
|
||||
return hidden_states
|
||||
|
||||
|
||||
class Downsample3D(nn.Module):
|
||||
def __init__(
|
||||
self, channels, use_conv=False, out_channels=None, padding=1, name="conv"
|
||||
):
|
||||
super().__init__()
|
||||
self.channels = channels
|
||||
self.out_channels = out_channels or channels
|
||||
self.use_conv = use_conv
|
||||
self.padding = padding
|
||||
stride = 2
|
||||
self.name = name
|
||||
|
||||
if use_conv:
|
||||
self.conv = InflatedConv3d(
|
||||
self.channels, self.out_channels, 3, stride=stride, padding=padding
|
||||
)
|
||||
else:
|
||||
raise NotImplementedError
|
||||
|
||||
def forward(self, hidden_states):
|
||||
assert hidden_states.shape[1] == self.channels
|
||||
if self.use_conv and self.padding == 0:
|
||||
raise NotImplementedError
|
||||
|
||||
assert hidden_states.shape[1] == self.channels
|
||||
hidden_states = self.conv(hidden_states)
|
||||
|
||||
return hidden_states
|
||||
|
||||
|
||||
class ResnetBlock3D(nn.Module):
|
||||
def __init__(
|
||||
self,
|
||||
*,
|
||||
in_channels,
|
||||
out_channels=None,
|
||||
conv_shortcut=False,
|
||||
dropout=0.0,
|
||||
temb_channels=512,
|
||||
groups=32,
|
||||
groups_out=None,
|
||||
pre_norm=True,
|
||||
eps=1e-6,
|
||||
non_linearity="swish",
|
||||
time_embedding_norm="default",
|
||||
output_scale_factor=1.0,
|
||||
use_in_shortcut=None,
|
||||
use_inflated_groupnorm=None,
|
||||
):
|
||||
super().__init__()
|
||||
self.pre_norm = pre_norm
|
||||
self.pre_norm = True
|
||||
self.in_channels = in_channels
|
||||
out_channels = in_channels if out_channels is None else out_channels
|
||||
self.out_channels = out_channels
|
||||
self.use_conv_shortcut = conv_shortcut
|
||||
self.time_embedding_norm = time_embedding_norm
|
||||
self.output_scale_factor = output_scale_factor
|
||||
|
||||
if groups_out is None:
|
||||
groups_out = groups
|
||||
|
||||
assert use_inflated_groupnorm != None
|
||||
if use_inflated_groupnorm:
|
||||
self.norm1 = InflatedGroupNorm(
|
||||
num_groups=groups, num_channels=in_channels, eps=eps, affine=True
|
||||
)
|
||||
else:
|
||||
self.norm1 = torch.nn.GroupNorm(
|
||||
num_groups=groups, num_channels=in_channels, eps=eps, affine=True
|
||||
)
|
||||
|
||||
self.conv1 = InflatedConv3d(
|
||||
in_channels, out_channels, kernel_size=3, stride=1, padding=1
|
||||
)
|
||||
|
||||
if temb_channels is not None:
|
||||
if self.time_embedding_norm == "default":
|
||||
time_emb_proj_out_channels = out_channels
|
||||
elif self.time_embedding_norm == "scale_shift":
|
||||
time_emb_proj_out_channels = out_channels * 2
|
||||
else:
|
||||
raise ValueError(
|
||||
f"unknown time_embedding_norm : {self.time_embedding_norm} "
|
||||
)
|
||||
|
||||
self.time_emb_proj = torch.nn.Linear(
|
||||
temb_channels, time_emb_proj_out_channels
|
||||
)
|
||||
else:
|
||||
self.time_emb_proj = None
|
||||
|
||||
if use_inflated_groupnorm:
|
||||
self.norm2 = InflatedGroupNorm(
|
||||
num_groups=groups_out, num_channels=out_channels, eps=eps, affine=True
|
||||
)
|
||||
else:
|
||||
self.norm2 = torch.nn.GroupNorm(
|
||||
num_groups=groups_out, num_channels=out_channels, eps=eps, affine=True
|
||||
)
|
||||
self.dropout = torch.nn.Dropout(dropout)
|
||||
self.conv2 = InflatedConv3d(
|
||||
out_channels, out_channels, kernel_size=3, stride=1, padding=1
|
||||
)
|
||||
|
||||
if non_linearity == "swish":
|
||||
self.nonlinearity = lambda x: F.silu(x)
|
||||
elif non_linearity == "mish":
|
||||
self.nonlinearity = Mish()
|
||||
elif non_linearity == "silu":
|
||||
self.nonlinearity = nn.SiLU()
|
||||
|
||||
self.use_in_shortcut = (
|
||||
self.in_channels != self.out_channels
|
||||
if use_in_shortcut is None
|
||||
else use_in_shortcut
|
||||
)
|
||||
|
||||
self.conv_shortcut = None
|
||||
if self.use_in_shortcut:
|
||||
self.conv_shortcut = InflatedConv3d(
|
||||
in_channels, out_channels, kernel_size=1, stride=1, padding=0
|
||||
)
|
||||
|
||||
def forward(self, input_tensor, temb):
|
||||
hidden_states = input_tensor
|
||||
|
||||
hidden_states = self.norm1(hidden_states)
|
||||
hidden_states = self.nonlinearity(hidden_states)
|
||||
|
||||
hidden_states = self.conv1(hidden_states)
|
||||
|
||||
if temb is not None:
|
||||
temb = self.time_emb_proj(self.nonlinearity(temb))[:, :, None, None, None]
|
||||
|
||||
if temb is not None and self.time_embedding_norm == "default":
|
||||
hidden_states = hidden_states + temb
|
||||
|
||||
hidden_states = self.norm2(hidden_states)
|
||||
|
||||
if temb is not None and self.time_embedding_norm == "scale_shift":
|
||||
scale, shift = torch.chunk(temb, 2, dim=1)
|
||||
hidden_states = hidden_states * (1 + scale) + shift
|
||||
|
||||
hidden_states = self.nonlinearity(hidden_states)
|
||||
|
||||
hidden_states = self.dropout(hidden_states)
|
||||
hidden_states = self.conv2(hidden_states)
|
||||
|
||||
if self.conv_shortcut is not None:
|
||||
input_tensor = self.conv_shortcut(input_tensor)
|
||||
|
||||
output_tensor = (input_tensor + hidden_states) / self.output_scale_factor
|
||||
|
||||
return output_tensor
|
||||
|
||||
|
||||
class Mish(torch.nn.Module):
|
||||
def forward(self, hidden_states):
|
||||
return hidden_states * torch.tanh(torch.nn.functional.softplus(hidden_states))
|
||||
@@ -0,0 +1,399 @@
|
||||
# Adapted from https://github.com/huggingface/diffusers/blob/main/src/diffusers/models/transformer_2d.py
|
||||
from dataclasses import dataclass
|
||||
from typing import Any, Dict, Optional
|
||||
|
||||
import torch
|
||||
from diffusers.configuration_utils import ConfigMixin, register_to_config
|
||||
try:
|
||||
from diffusers.models.embeddings import CaptionProjection
|
||||
except ImportError:
|
||||
from diffusers.models.embeddings import PixArtAlphaTextProjection as CaptionProjection
|
||||
from diffusers.models.lora import LoRACompatibleConv, LoRACompatibleLinear
|
||||
from diffusers.models.modeling_utils import ModelMixin
|
||||
from diffusers.models.normalization import AdaLayerNormSingle
|
||||
from diffusers.utils import USE_PEFT_BACKEND, BaseOutput, deprecate, is_torch_version
|
||||
from torch import nn
|
||||
|
||||
from .attention import BasicTransformerBlock
|
||||
|
||||
|
||||
@dataclass
|
||||
class Transformer2DModelOutput(BaseOutput):
|
||||
"""
|
||||
The output of [`Transformer2DModel`].
|
||||
|
||||
Args:
|
||||
sample (`torch.FloatTensor` of shape `(batch_size, num_channels, height, width)` or `(batch size, num_vector_embeds - 1, num_latent_pixels)` if [`Transformer2DModel`] is discrete):
|
||||
The hidden states output conditioned on the `encoder_hidden_states` input. If discrete, returns probability
|
||||
distributions for the unnoised latent pixels.
|
||||
"""
|
||||
|
||||
sample: torch.FloatTensor
|
||||
ref_feature: torch.FloatTensor
|
||||
|
||||
|
||||
class Transformer2DModel(ModelMixin, ConfigMixin):
|
||||
"""
|
||||
A 2D Transformer model for image-like data.
|
||||
|
||||
Parameters:
|
||||
num_attention_heads (`int`, *optional*, defaults to 16): The number of heads to use for multi-head attention.
|
||||
attention_head_dim (`int`, *optional*, defaults to 88): The number of channels in each head.
|
||||
in_channels (`int`, *optional*):
|
||||
The number of channels in the input and output (specify if the input is **continuous**).
|
||||
num_layers (`int`, *optional*, defaults to 1): The number of layers of Transformer blocks to use.
|
||||
dropout (`float`, *optional*, defaults to 0.0): The dropout probability to use.
|
||||
cross_attention_dim (`int`, *optional*): The number of `encoder_hidden_states` dimensions to use.
|
||||
sample_size (`int`, *optional*): The width of the latent images (specify if the input is **discrete**).
|
||||
This is fixed during training since it is used to learn a number of position embeddings.
|
||||
num_vector_embeds (`int`, *optional*):
|
||||
The number of classes of the vector embeddings of the latent pixels (specify if the input is **discrete**).
|
||||
Includes the class for the masked latent pixel.
|
||||
activation_fn (`str`, *optional*, defaults to `"geglu"`): Activation function to use in feed-forward.
|
||||
num_embeds_ada_norm ( `int`, *optional*):
|
||||
The number of diffusion steps used during training. Pass if at least one of the norm_layers is
|
||||
`AdaLayerNorm`. This is fixed during training since it is used to learn a number of embeddings that are
|
||||
added to the hidden states.
|
||||
|
||||
During inference, you can denoise for up to but not more steps than `num_embeds_ada_norm`.
|
||||
attention_bias (`bool`, *optional*):
|
||||
Configure if the `TransformerBlocks` attention should contain a bias parameter.
|
||||
"""
|
||||
|
||||
_supports_gradient_checkpointing = True
|
||||
|
||||
@register_to_config
|
||||
def __init__(
|
||||
self,
|
||||
num_attention_heads: int = 16,
|
||||
attention_head_dim: int = 88,
|
||||
in_channels: Optional[int] = None,
|
||||
out_channels: Optional[int] = None,
|
||||
num_layers: int = 1,
|
||||
dropout: float = 0.0,
|
||||
norm_num_groups: int = 32,
|
||||
cross_attention_dim: Optional[int] = None,
|
||||
attention_bias: bool = False,
|
||||
sample_size: Optional[int] = None,
|
||||
num_vector_embeds: Optional[int] = None,
|
||||
patch_size: Optional[int] = None,
|
||||
activation_fn: str = "geglu",
|
||||
num_embeds_ada_norm: Optional[int] = None,
|
||||
use_linear_projection: bool = False,
|
||||
only_cross_attention: bool = False,
|
||||
double_self_attention: bool = False,
|
||||
upcast_attention: bool = False,
|
||||
norm_type: str = "layer_norm",
|
||||
norm_elementwise_affine: bool = True,
|
||||
norm_eps: float = 1e-5,
|
||||
attention_type: str = "default",
|
||||
caption_channels: int = None,
|
||||
):
|
||||
super().__init__()
|
||||
self.use_linear_projection = use_linear_projection
|
||||
self.num_attention_heads = num_attention_heads
|
||||
self.attention_head_dim = attention_head_dim
|
||||
inner_dim = num_attention_heads * attention_head_dim
|
||||
|
||||
conv_cls = nn.Conv2d if USE_PEFT_BACKEND else LoRACompatibleConv
|
||||
linear_cls = nn.Linear if USE_PEFT_BACKEND else LoRACompatibleLinear
|
||||
|
||||
# 1. Transformer2DModel can process both standard continuous images of shape `(batch_size, num_channels, width, height)` as well as quantized image embeddings of shape `(batch_size, num_image_vectors)`
|
||||
# Define whether input is continuous or discrete depending on configuration
|
||||
self.is_input_continuous = (in_channels is not None) and (patch_size is None)
|
||||
self.is_input_vectorized = num_vector_embeds is not None
|
||||
self.is_input_patches = in_channels is not None and patch_size is not None
|
||||
|
||||
if norm_type == "layer_norm" and num_embeds_ada_norm is not None:
|
||||
deprecation_message = (
|
||||
f"The configuration file of this model: {self.__class__} is outdated. `norm_type` is either not set or"
|
||||
" incorrectly set to `'layer_norm'`.Make sure to set `norm_type` to `'ada_norm'` in the config."
|
||||
" Please make sure to update the config accordingly as leaving `norm_type` might led to incorrect"
|
||||
" results in future versions. If you have downloaded this checkpoint from the Hugging Face Hub, it"
|
||||
" would be very nice if you could open a Pull request for the `transformer/config.json` file"
|
||||
)
|
||||
deprecate(
|
||||
"norm_type!=num_embeds_ada_norm",
|
||||
"1.0.0",
|
||||
deprecation_message,
|
||||
standard_warn=False,
|
||||
)
|
||||
norm_type = "ada_norm"
|
||||
|
||||
if self.is_input_continuous and self.is_input_vectorized:
|
||||
raise ValueError(
|
||||
f"Cannot define both `in_channels`: {in_channels} and `num_vector_embeds`: {num_vector_embeds}. Make"
|
||||
" sure that either `in_channels` or `num_vector_embeds` is None."
|
||||
)
|
||||
elif self.is_input_vectorized and self.is_input_patches:
|
||||
raise ValueError(
|
||||
f"Cannot define both `num_vector_embeds`: {num_vector_embeds} and `patch_size`: {patch_size}. Make"
|
||||
" sure that either `num_vector_embeds` or `num_patches` is None."
|
||||
)
|
||||
elif (
|
||||
not self.is_input_continuous
|
||||
and not self.is_input_vectorized
|
||||
and not self.is_input_patches
|
||||
):
|
||||
raise ValueError(
|
||||
f"Has to define `in_channels`: {in_channels}, `num_vector_embeds`: {num_vector_embeds}, or patch_size:"
|
||||
f" {patch_size}. Make sure that `in_channels`, `num_vector_embeds` or `num_patches` is not None."
|
||||
)
|
||||
|
||||
# 2. Define input layers
|
||||
self.in_channels = in_channels
|
||||
|
||||
self.norm = torch.nn.GroupNorm(
|
||||
num_groups=norm_num_groups,
|
||||
num_channels=in_channels,
|
||||
eps=1e-6,
|
||||
affine=True,
|
||||
)
|
||||
if use_linear_projection:
|
||||
self.proj_in = linear_cls(in_channels, inner_dim)
|
||||
else:
|
||||
self.proj_in = conv_cls(
|
||||
in_channels, inner_dim, kernel_size=1, stride=1, padding=0
|
||||
)
|
||||
|
||||
# 3. Define transformers blocks
|
||||
self.transformer_blocks = nn.ModuleList(
|
||||
[
|
||||
BasicTransformerBlock(
|
||||
inner_dim,
|
||||
num_attention_heads,
|
||||
attention_head_dim,
|
||||
dropout=dropout,
|
||||
cross_attention_dim=cross_attention_dim,
|
||||
activation_fn=activation_fn,
|
||||
num_embeds_ada_norm=num_embeds_ada_norm,
|
||||
attention_bias=attention_bias,
|
||||
only_cross_attention=only_cross_attention,
|
||||
double_self_attention=double_self_attention,
|
||||
upcast_attention=upcast_attention,
|
||||
norm_type=norm_type,
|
||||
norm_elementwise_affine=norm_elementwise_affine,
|
||||
norm_eps=norm_eps,
|
||||
attention_type=attention_type,
|
||||
)
|
||||
for d in range(num_layers)
|
||||
]
|
||||
)
|
||||
|
||||
# 4. Define output layers
|
||||
self.out_channels = in_channels if out_channels is None else out_channels
|
||||
# TODO: should use out_channels for continuous projections
|
||||
if use_linear_projection:
|
||||
self.proj_out = linear_cls(inner_dim, in_channels)
|
||||
else:
|
||||
self.proj_out = conv_cls(
|
||||
inner_dim, in_channels, kernel_size=1, stride=1, padding=0
|
||||
)
|
||||
|
||||
# 5. PixArt-Alpha blocks.
|
||||
self.adaln_single = None
|
||||
self.use_additional_conditions = False
|
||||
if norm_type == "ada_norm_single":
|
||||
self.use_additional_conditions = self.config.sample_size == 128
|
||||
# TODO(Sayak, PVP) clean this, for now we use sample size to determine whether to use
|
||||
# additional conditions until we find better name
|
||||
self.adaln_single = AdaLayerNormSingle(
|
||||
inner_dim, use_additional_conditions=self.use_additional_conditions
|
||||
)
|
||||
|
||||
self.caption_projection = None
|
||||
if caption_channels is not None:
|
||||
self.caption_projection = CaptionProjection(
|
||||
in_features=caption_channels, hidden_size=inner_dim
|
||||
)
|
||||
|
||||
self.gradient_checkpointing = False
|
||||
|
||||
def _set_gradient_checkpointing(self, module, value=False):
|
||||
if hasattr(module, "gradient_checkpointing"):
|
||||
module.gradient_checkpointing = value
|
||||
|
||||
def forward(
|
||||
self,
|
||||
hidden_states: torch.Tensor,
|
||||
encoder_hidden_states: Optional[torch.Tensor] = None,
|
||||
timestep: Optional[torch.LongTensor] = None,
|
||||
added_cond_kwargs: Dict[str, torch.Tensor] = None,
|
||||
class_labels: Optional[torch.LongTensor] = None,
|
||||
cross_attention_kwargs: Dict[str, Any] = None,
|
||||
attention_mask: Optional[torch.Tensor] = None,
|
||||
encoder_attention_mask: Optional[torch.Tensor] = None,
|
||||
return_dict: bool = True,
|
||||
):
|
||||
"""
|
||||
The [`Transformer2DModel`] forward method.
|
||||
|
||||
Args:
|
||||
hidden_states (`torch.LongTensor` of shape `(batch size, num latent pixels)` if discrete, `torch.FloatTensor` of shape `(batch size, channel, height, width)` if continuous):
|
||||
Input `hidden_states`.
|
||||
encoder_hidden_states ( `torch.FloatTensor` of shape `(batch size, sequence len, embed dims)`, *optional*):
|
||||
Conditional embeddings for cross attention layer. If not given, cross-attention defaults to
|
||||
self-attention.
|
||||
timestep ( `torch.LongTensor`, *optional*):
|
||||
Used to indicate denoising step. Optional timestep to be applied as an embedding in `AdaLayerNorm`.
|
||||
class_labels ( `torch.LongTensor` of shape `(batch size, num classes)`, *optional*):
|
||||
Used to indicate class labels conditioning. Optional class labels to be applied as an embedding in
|
||||
`AdaLayerZeroNorm`.
|
||||
cross_attention_kwargs ( `Dict[str, Any]`, *optional*):
|
||||
A kwargs dictionary that if specified is passed along to the `AttentionProcessor` as defined under
|
||||
`self.processor` in
|
||||
[diffusers.models.attention_processor](https://github.com/huggingface/diffusers/blob/main/src/diffusers/models/attention_processor.py).
|
||||
attention_mask ( `torch.Tensor`, *optional*):
|
||||
An attention mask of shape `(batch, key_tokens)` is applied to `encoder_hidden_states`. If `1` the mask
|
||||
is kept, otherwise if `0` it is discarded. Mask will be converted into a bias, which adds large
|
||||
negative values to the attention scores corresponding to "discard" tokens.
|
||||
encoder_attention_mask ( `torch.Tensor`, *optional*):
|
||||
Cross-attention mask applied to `encoder_hidden_states`. Two formats supported:
|
||||
|
||||
* Mask `(batch, sequence_length)` True = keep, False = discard.
|
||||
* Bias `(batch, 1, sequence_length)` 0 = keep, -10000 = discard.
|
||||
|
||||
If `ndim == 2`: will be interpreted as a mask, then converted into a bias consistent with the format
|
||||
above. This bias will be added to the cross-attention scores.
|
||||
return_dict (`bool`, *optional*, defaults to `True`):
|
||||
Whether or not to return a [`~models.unet_2d_condition.UNet2DConditionOutput`] instead of a plain
|
||||
tuple.
|
||||
|
||||
Returns:
|
||||
If `return_dict` is True, an [`~models.transformer_2d.Transformer2DModelOutput`] is returned, otherwise a
|
||||
`tuple` where the first element is the sample tensor.
|
||||
"""
|
||||
# ensure attention_mask is a bias, and give it a singleton query_tokens dimension.
|
||||
# we may have done this conversion already, e.g. if we came here via UNet2DConditionModel#forward.
|
||||
# we can tell by counting dims; if ndim == 2: it's a mask rather than a bias.
|
||||
# expects mask of shape:
|
||||
# [batch, key_tokens]
|
||||
# adds singleton query_tokens dimension:
|
||||
# [batch, 1, key_tokens]
|
||||
# this helps to broadcast it as a bias over attention scores, which will be in one of the following shapes:
|
||||
# [batch, heads, query_tokens, key_tokens] (e.g. torch sdp attn)
|
||||
# [batch * heads, query_tokens, key_tokens] (e.g. xformers or classic attn)
|
||||
if attention_mask is not None and attention_mask.ndim == 2:
|
||||
# assume that mask is expressed as:
|
||||
# (1 = keep, 0 = discard)
|
||||
# convert mask into a bias that can be added to attention scores:
|
||||
# (keep = +0, discard = -10000.0)
|
||||
attention_mask = (1 - attention_mask.to(hidden_states.dtype)) * -10000.0
|
||||
attention_mask = attention_mask.unsqueeze(1)
|
||||
|
||||
# convert encoder_attention_mask to a bias the same way we do for attention_mask
|
||||
if encoder_attention_mask is not None and encoder_attention_mask.ndim == 2:
|
||||
encoder_attention_mask = (
|
||||
1 - encoder_attention_mask.to(hidden_states.dtype)
|
||||
) * -10000.0
|
||||
encoder_attention_mask = encoder_attention_mask.unsqueeze(1)
|
||||
|
||||
# Retrieve lora scale.
|
||||
lora_scale = (
|
||||
cross_attention_kwargs.get("scale", 1.0)
|
||||
if cross_attention_kwargs is not None
|
||||
else 1.0
|
||||
)
|
||||
|
||||
# 1. Input
|
||||
batch, _, height, width = hidden_states.shape
|
||||
residual = hidden_states
|
||||
|
||||
hidden_states = self.norm(hidden_states)
|
||||
if not self.use_linear_projection:
|
||||
hidden_states = (
|
||||
self.proj_in(hidden_states, scale=lora_scale)
|
||||
if not USE_PEFT_BACKEND
|
||||
else self.proj_in(hidden_states)
|
||||
)
|
||||
inner_dim = hidden_states.shape[1]
|
||||
hidden_states = hidden_states.permute(0, 2, 3, 1).reshape(
|
||||
batch, height * width, inner_dim
|
||||
)
|
||||
else:
|
||||
inner_dim = hidden_states.shape[1]
|
||||
hidden_states = hidden_states.permute(0, 2, 3, 1).reshape(
|
||||
batch, height * width, inner_dim
|
||||
)
|
||||
hidden_states = (
|
||||
self.proj_in(hidden_states, scale=lora_scale)
|
||||
if not USE_PEFT_BACKEND
|
||||
else self.proj_in(hidden_states)
|
||||
)
|
||||
|
||||
# 2. Blocks
|
||||
if self.caption_projection is not None:
|
||||
batch_size = hidden_states.shape[0]
|
||||
encoder_hidden_states = self.caption_projection(encoder_hidden_states)
|
||||
encoder_hidden_states = encoder_hidden_states.view(
|
||||
batch_size, -1, hidden_states.shape[-1]
|
||||
)
|
||||
|
||||
ref_feature = hidden_states.reshape(batch, height, width, inner_dim)
|
||||
for block in self.transformer_blocks:
|
||||
if self.training and self.gradient_checkpointing:
|
||||
|
||||
def create_custom_forward(module, return_dict=None):
|
||||
def custom_forward(*inputs):
|
||||
if return_dict is not None:
|
||||
return module(*inputs, return_dict=return_dict)
|
||||
else:
|
||||
return module(*inputs)
|
||||
|
||||
return custom_forward
|
||||
|
||||
ckpt_kwargs: Dict[str, Any] = (
|
||||
{"use_reentrant": False} if is_torch_version(">=", "1.11.0") else {}
|
||||
)
|
||||
hidden_states = torch.utils.checkpoint.checkpoint(
|
||||
create_custom_forward(block),
|
||||
hidden_states,
|
||||
attention_mask,
|
||||
encoder_hidden_states,
|
||||
encoder_attention_mask,
|
||||
timestep,
|
||||
cross_attention_kwargs,
|
||||
class_labels,
|
||||
**ckpt_kwargs,
|
||||
)
|
||||
else:
|
||||
hidden_states = block(
|
||||
hidden_states,
|
||||
attention_mask=attention_mask,
|
||||
encoder_hidden_states=encoder_hidden_states,
|
||||
encoder_attention_mask=encoder_attention_mask,
|
||||
timestep=timestep,
|
||||
cross_attention_kwargs=cross_attention_kwargs,
|
||||
class_labels=class_labels,
|
||||
)
|
||||
|
||||
# 3. Output
|
||||
if self.is_input_continuous:
|
||||
if not self.use_linear_projection:
|
||||
hidden_states = (
|
||||
hidden_states.reshape(batch, height, width, inner_dim)
|
||||
.permute(0, 3, 1, 2)
|
||||
.contiguous()
|
||||
)
|
||||
hidden_states = (
|
||||
self.proj_out(hidden_states, scale=lora_scale)
|
||||
if not USE_PEFT_BACKEND
|
||||
else self.proj_out(hidden_states)
|
||||
)
|
||||
else:
|
||||
hidden_states = (
|
||||
self.proj_out(hidden_states, scale=lora_scale)
|
||||
if not USE_PEFT_BACKEND
|
||||
else self.proj_out(hidden_states)
|
||||
)
|
||||
hidden_states = (
|
||||
hidden_states.reshape(batch, height, width, inner_dim)
|
||||
.permute(0, 3, 1, 2)
|
||||
.contiguous()
|
||||
)
|
||||
|
||||
output = hidden_states + residual
|
||||
if not return_dict:
|
||||
return (output, ref_feature)
|
||||
|
||||
return Transformer2DModelOutput(sample=output, ref_feature=ref_feature)
|
||||
@@ -0,0 +1,175 @@
|
||||
from dataclasses import dataclass
|
||||
from typing import Optional
|
||||
|
||||
import torch
|
||||
from diffusers.configuration_utils import ConfigMixin, register_to_config
|
||||
from diffusers.models import ModelMixin
|
||||
from diffusers.utils import BaseOutput
|
||||
from diffusers.utils.import_utils import is_xformers_available
|
||||
from einops import rearrange, repeat
|
||||
from torch import nn
|
||||
|
||||
from .attention import TemporalBasicTransformerBlock
|
||||
|
||||
|
||||
@dataclass
|
||||
class Transformer3DModelOutput(BaseOutput):
|
||||
sample: torch.FloatTensor
|
||||
|
||||
|
||||
if is_xformers_available():
|
||||
import xformers
|
||||
import xformers.ops
|
||||
else:
|
||||
xformers = None
|
||||
|
||||
|
||||
class Transformer3DModel(ModelMixin, ConfigMixin):
|
||||
_supports_gradient_checkpointing = True
|
||||
|
||||
@register_to_config
|
||||
def __init__(
|
||||
self,
|
||||
num_attention_heads: int = 16,
|
||||
attention_head_dim: int = 88,
|
||||
in_channels: Optional[int] = None,
|
||||
num_layers: int = 1,
|
||||
dropout: float = 0.0,
|
||||
norm_num_groups: int = 32,
|
||||
cross_attention_dim: Optional[int] = None,
|
||||
attention_bias: bool = False,
|
||||
activation_fn: str = "geglu",
|
||||
num_embeds_ada_norm: Optional[int] = None,
|
||||
use_linear_projection: bool = False,
|
||||
only_cross_attention: bool = False,
|
||||
upcast_attention: bool = False,
|
||||
unet_use_cross_frame_attention=None,
|
||||
unet_use_temporal_attention=None,
|
||||
):
|
||||
super().__init__()
|
||||
self.use_linear_projection = use_linear_projection
|
||||
self.num_attention_heads = num_attention_heads
|
||||
self.attention_head_dim = attention_head_dim
|
||||
inner_dim = num_attention_heads * attention_head_dim
|
||||
|
||||
# Define input layers
|
||||
self.in_channels = in_channels
|
||||
|
||||
self.norm = torch.nn.GroupNorm(
|
||||
num_groups=norm_num_groups, num_channels=in_channels, eps=1e-6, affine=True
|
||||
)
|
||||
if use_linear_projection:
|
||||
self.proj_in = nn.Linear(in_channels, inner_dim)
|
||||
else:
|
||||
self.proj_in = nn.Conv2d(
|
||||
in_channels, inner_dim, kernel_size=1, stride=1, padding=0
|
||||
)
|
||||
|
||||
# Define transformers blocks
|
||||
self.transformer_blocks = nn.ModuleList(
|
||||
[
|
||||
TemporalBasicTransformerBlock(
|
||||
inner_dim,
|
||||
num_attention_heads,
|
||||
attention_head_dim,
|
||||
dropout=dropout,
|
||||
cross_attention_dim=cross_attention_dim,
|
||||
activation_fn=activation_fn,
|
||||
num_embeds_ada_norm=num_embeds_ada_norm,
|
||||
attention_bias=attention_bias,
|
||||
only_cross_attention=only_cross_attention,
|
||||
upcast_attention=upcast_attention,
|
||||
unet_use_cross_frame_attention=unet_use_cross_frame_attention,
|
||||
unet_use_temporal_attention=unet_use_temporal_attention,
|
||||
)
|
||||
for d in range(num_layers)
|
||||
]
|
||||
)
|
||||
|
||||
# 4. Define output layers
|
||||
if use_linear_projection:
|
||||
self.proj_out = nn.Linear(in_channels, inner_dim)
|
||||
else:
|
||||
self.proj_out = nn.Conv2d(
|
||||
inner_dim, in_channels, kernel_size=1, stride=1, padding=0
|
||||
)
|
||||
|
||||
self.gradient_checkpointing = False
|
||||
|
||||
def _set_gradient_checkpointing(self, module, value=False):
|
||||
if hasattr(module, "gradient_checkpointing"):
|
||||
module.gradient_checkpointing = value
|
||||
|
||||
def forward(
|
||||
self,
|
||||
hidden_states,
|
||||
encoder_hidden_states=None,
|
||||
audio_cond_fea=None,
|
||||
timestep=None,
|
||||
return_dict: bool = True,
|
||||
):
|
||||
# Input
|
||||
assert (
|
||||
hidden_states.dim() == 5
|
||||
), f"Expected hidden_states to have ndim=5, but got ndim={hidden_states.dim()}."
|
||||
video_length = hidden_states.shape[2]
|
||||
hidden_states = rearrange(hidden_states, "b c f h w -> (b f) c h w")
|
||||
if audio_cond_fea.shape[0] != hidden_states.shape[0]:
|
||||
if len(audio_cond_fea.shape) == 3:
|
||||
audio_cond_fea = rearrange(
|
||||
audio_cond_fea, "b f c -> (b f) 1 c"
|
||||
)
|
||||
elif len(audio_cond_fea.shape) == 4:
|
||||
audio_cond_fea = rearrange(
|
||||
audio_cond_fea, "b f n c -> (b f) n c"
|
||||
)
|
||||
batch, channel, height, weight = hidden_states.shape
|
||||
residual = hidden_states
|
||||
|
||||
hidden_states = self.norm(hidden_states)
|
||||
if not self.use_linear_projection:
|
||||
hidden_states = self.proj_in(hidden_states)
|
||||
inner_dim = hidden_states.shape[1]
|
||||
hidden_states = hidden_states.permute(0, 2, 3, 1).reshape(
|
||||
batch, height * weight, inner_dim
|
||||
)
|
||||
else:
|
||||
inner_dim = hidden_states.shape[1]
|
||||
hidden_states = hidden_states.permute(0, 2, 3, 1).reshape(
|
||||
batch, height * weight, inner_dim
|
||||
)
|
||||
hidden_states = self.proj_in(hidden_states)
|
||||
|
||||
# Blocks
|
||||
for i, block in enumerate(self.transformer_blocks):
|
||||
hidden_states = block(
|
||||
hidden_states,
|
||||
encoder_hidden_states=encoder_hidden_states,
|
||||
audio_cond_fea=audio_cond_fea,
|
||||
timestep=timestep,
|
||||
video_length=video_length,
|
||||
)
|
||||
|
||||
# Output
|
||||
if not self.use_linear_projection:
|
||||
hidden_states = (
|
||||
hidden_states.reshape(batch, height, weight, inner_dim)
|
||||
.permute(0, 3, 1, 2)
|
||||
.contiguous()
|
||||
)
|
||||
hidden_states = self.proj_out(hidden_states)
|
||||
else:
|
||||
hidden_states = self.proj_out(hidden_states)
|
||||
hidden_states = (
|
||||
hidden_states.reshape(batch, height, weight, inner_dim)
|
||||
.permute(0, 3, 1, 2)
|
||||
.contiguous()
|
||||
)
|
||||
|
||||
output = hidden_states + residual
|
||||
|
||||
output = rearrange(output, "(b f) c h w -> b c f h w", f=video_length)
|
||||
if not return_dict:
|
||||
return (output,)
|
||||
|
||||
return Transformer3DModelOutput(sample=output)
|
||||
File diff suppressed because it is too large
Load Diff
File diff suppressed because it is too large
Load Diff
@@ -0,0 +1,668 @@
|
||||
# Adapted from https://github.com/guoyww/AnimateDiff/blob/main/animatediff/models/unet_blocks.py
|
||||
|
||||
from collections import OrderedDict
|
||||
from dataclasses import dataclass
|
||||
from os import PathLike
|
||||
from pathlib import Path
|
||||
from typing import Dict, List, Optional, Tuple, Union
|
||||
|
||||
import torch
|
||||
import torch.nn as nn
|
||||
import torch.utils.checkpoint
|
||||
from diffusers.configuration_utils import ConfigMixin, register_to_config
|
||||
from diffusers.models.attention_processor import AttentionProcessor
|
||||
from diffusers.models.embeddings import TimestepEmbedding, Timesteps
|
||||
from diffusers.models.modeling_utils import ModelMixin
|
||||
from diffusers.utils import SAFETENSORS_WEIGHTS_NAME, WEIGHTS_NAME, BaseOutput, logging
|
||||
from safetensors.torch import load_file
|
||||
|
||||
from .resnet import InflatedConv3d, InflatedGroupNorm
|
||||
from .unet_3d_blocks import UNetMidBlock3DCrossAttn, get_down_block, get_up_block
|
||||
|
||||
logger = logging.get_logger(__name__) # pylint: disable=invalid-name
|
||||
|
||||
|
||||
@dataclass
|
||||
class UNet3DConditionOutput(BaseOutput):
|
||||
sample: torch.FloatTensor
|
||||
|
||||
|
||||
class UNet3DConditionModel(ModelMixin, ConfigMixin):
|
||||
_supports_gradient_checkpointing = True
|
||||
|
||||
@register_to_config
|
||||
def __init__(
|
||||
self,
|
||||
sample_size: Optional[int] = None,
|
||||
in_channels: int = 4,
|
||||
out_channels: int = 4,
|
||||
center_input_sample: bool = False,
|
||||
flip_sin_to_cos: bool = True,
|
||||
freq_shift: int = 0,
|
||||
down_block_types: Tuple[str] = (
|
||||
"CrossAttnDownBlock3D",
|
||||
"CrossAttnDownBlock3D",
|
||||
"CrossAttnDownBlock3D",
|
||||
"DownBlock3D",
|
||||
),
|
||||
mid_block_type: str = "UNetMidBlock3DCrossAttn",
|
||||
up_block_types: Tuple[str] = (
|
||||
"UpBlock3D",
|
||||
"CrossAttnUpBlock3D",
|
||||
"CrossAttnUpBlock3D",
|
||||
"CrossAttnUpBlock3D",
|
||||
),
|
||||
only_cross_attention: Union[bool, Tuple[bool]] = False,
|
||||
block_out_channels: Tuple[int] = (320, 640, 1280, 1280),
|
||||
layers_per_block: int = 2,
|
||||
downsample_padding: int = 1,
|
||||
mid_block_scale_factor: float = 1,
|
||||
act_fn: str = "silu",
|
||||
norm_num_groups: int = 32,
|
||||
norm_eps: float = 1e-5,
|
||||
cross_attention_dim: int = 1280,
|
||||
attention_head_dim: Union[int, Tuple[int]] = 8,
|
||||
dual_cross_attention: bool = False,
|
||||
use_linear_projection: bool = False,
|
||||
class_embed_type: Optional[str] = None,
|
||||
num_class_embeds: Optional[int] = None,
|
||||
upcast_attention: bool = False,
|
||||
resnet_time_scale_shift: str = "default",
|
||||
use_inflated_groupnorm=False,
|
||||
# Additional
|
||||
use_motion_module=False,
|
||||
motion_module_resolutions=(1, 2, 4, 8),
|
||||
motion_module_mid_block=False,
|
||||
motion_module_decoder_only=False,
|
||||
motion_module_type=None,
|
||||
motion_module_kwargs={},
|
||||
unet_use_cross_frame_attention=None,
|
||||
unet_use_temporal_attention=None,
|
||||
):
|
||||
super().__init__()
|
||||
|
||||
self.sample_size = sample_size
|
||||
time_embed_dim = block_out_channels[0] * 4
|
||||
|
||||
# input
|
||||
self.conv_in = InflatedConv3d(
|
||||
in_channels, block_out_channels[0], kernel_size=3, padding=(1, 1)
|
||||
)
|
||||
|
||||
# time
|
||||
self.time_proj = Timesteps(block_out_channels[0], flip_sin_to_cos, freq_shift)
|
||||
timestep_input_dim = block_out_channels[0]
|
||||
|
||||
self.time_embedding = TimestepEmbedding(timestep_input_dim, time_embed_dim)
|
||||
|
||||
# class embedding
|
||||
if class_embed_type is None and num_class_embeds is not None:
|
||||
self.class_embedding = nn.Embedding(num_class_embeds, time_embed_dim)
|
||||
elif class_embed_type == "timestep":
|
||||
self.class_embedding = TimestepEmbedding(timestep_input_dim, time_embed_dim)
|
||||
elif class_embed_type == "identity":
|
||||
self.class_embedding = nn.Identity(time_embed_dim, time_embed_dim)
|
||||
else:
|
||||
self.class_embedding = None
|
||||
|
||||
self.down_blocks = nn.ModuleList([])
|
||||
self.mid_block = None
|
||||
self.up_blocks = nn.ModuleList([])
|
||||
|
||||
if isinstance(only_cross_attention, bool):
|
||||
only_cross_attention = [only_cross_attention] * len(down_block_types)
|
||||
|
||||
if isinstance(attention_head_dim, int):
|
||||
attention_head_dim = (attention_head_dim,) * len(down_block_types)
|
||||
|
||||
# down
|
||||
output_channel = block_out_channels[0]
|
||||
for i, down_block_type in enumerate(down_block_types):
|
||||
res = 2**i
|
||||
input_channel = output_channel
|
||||
output_channel = block_out_channels[i]
|
||||
is_final_block = i == len(block_out_channels) - 1
|
||||
|
||||
down_block = get_down_block(
|
||||
down_block_type,
|
||||
num_layers=layers_per_block,
|
||||
in_channels=input_channel,
|
||||
out_channels=output_channel,
|
||||
temb_channels=time_embed_dim,
|
||||
add_downsample=not is_final_block,
|
||||
resnet_eps=norm_eps,
|
||||
resnet_act_fn=act_fn,
|
||||
resnet_groups=norm_num_groups,
|
||||
cross_attention_dim=cross_attention_dim,
|
||||
attn_num_head_channels=attention_head_dim[i],
|
||||
downsample_padding=downsample_padding,
|
||||
dual_cross_attention=dual_cross_attention,
|
||||
use_linear_projection=use_linear_projection,
|
||||
only_cross_attention=only_cross_attention[i],
|
||||
upcast_attention=upcast_attention,
|
||||
resnet_time_scale_shift=resnet_time_scale_shift,
|
||||
unet_use_cross_frame_attention=unet_use_cross_frame_attention,
|
||||
unet_use_temporal_attention=unet_use_temporal_attention,
|
||||
use_inflated_groupnorm=use_inflated_groupnorm,
|
||||
use_motion_module=use_motion_module
|
||||
and (res in motion_module_resolutions)
|
||||
and (not motion_module_decoder_only),
|
||||
motion_module_type=motion_module_type,
|
||||
motion_module_kwargs=motion_module_kwargs,
|
||||
)
|
||||
self.down_blocks.append(down_block)
|
||||
|
||||
# mid
|
||||
if mid_block_type == "UNetMidBlock3DCrossAttn":
|
||||
self.mid_block = UNetMidBlock3DCrossAttn(
|
||||
in_channels=block_out_channels[-1],
|
||||
temb_channels=time_embed_dim,
|
||||
resnet_eps=norm_eps,
|
||||
resnet_act_fn=act_fn,
|
||||
output_scale_factor=mid_block_scale_factor,
|
||||
resnet_time_scale_shift=resnet_time_scale_shift,
|
||||
cross_attention_dim=cross_attention_dim,
|
||||
attn_num_head_channels=attention_head_dim[-1],
|
||||
resnet_groups=norm_num_groups,
|
||||
dual_cross_attention=dual_cross_attention,
|
||||
use_linear_projection=use_linear_projection,
|
||||
upcast_attention=upcast_attention,
|
||||
unet_use_cross_frame_attention=unet_use_cross_frame_attention,
|
||||
unet_use_temporal_attention=unet_use_temporal_attention,
|
||||
use_inflated_groupnorm=use_inflated_groupnorm,
|
||||
use_motion_module=use_motion_module and motion_module_mid_block,
|
||||
motion_module_type=motion_module_type,
|
||||
motion_module_kwargs=motion_module_kwargs,
|
||||
)
|
||||
else:
|
||||
raise ValueError(f"unknown mid_block_type : {mid_block_type}")
|
||||
|
||||
# count how many layers upsample the videos
|
||||
self.num_upsamplers = 0
|
||||
|
||||
# up
|
||||
reversed_block_out_channels = list(reversed(block_out_channels))
|
||||
reversed_attention_head_dim = list(reversed(attention_head_dim))
|
||||
only_cross_attention = list(reversed(only_cross_attention))
|
||||
output_channel = reversed_block_out_channels[0]
|
||||
for i, up_block_type in enumerate(up_block_types):
|
||||
res = 2 ** (3 - i)
|
||||
is_final_block = i == len(block_out_channels) - 1
|
||||
|
||||
prev_output_channel = output_channel
|
||||
output_channel = reversed_block_out_channels[i]
|
||||
input_channel = reversed_block_out_channels[
|
||||
min(i + 1, len(block_out_channels) - 1)
|
||||
]
|
||||
|
||||
# add upsample block for all BUT final layer
|
||||
if not is_final_block:
|
||||
add_upsample = True
|
||||
self.num_upsamplers += 1
|
||||
else:
|
||||
add_upsample = False
|
||||
|
||||
up_block = get_up_block(
|
||||
up_block_type,
|
||||
num_layers=layers_per_block + 1,
|
||||
in_channels=input_channel,
|
||||
out_channels=output_channel,
|
||||
prev_output_channel=prev_output_channel,
|
||||
temb_channels=time_embed_dim,
|
||||
add_upsample=add_upsample,
|
||||
resnet_eps=norm_eps,
|
||||
resnet_act_fn=act_fn,
|
||||
resnet_groups=norm_num_groups,
|
||||
cross_attention_dim=cross_attention_dim,
|
||||
attn_num_head_channels=reversed_attention_head_dim[i],
|
||||
dual_cross_attention=dual_cross_attention,
|
||||
use_linear_projection=use_linear_projection,
|
||||
only_cross_attention=only_cross_attention[i],
|
||||
upcast_attention=upcast_attention,
|
||||
resnet_time_scale_shift=resnet_time_scale_shift,
|
||||
unet_use_cross_frame_attention=unet_use_cross_frame_attention,
|
||||
unet_use_temporal_attention=unet_use_temporal_attention,
|
||||
use_inflated_groupnorm=use_inflated_groupnorm,
|
||||
use_motion_module=use_motion_module
|
||||
and (res in motion_module_resolutions),
|
||||
motion_module_type=motion_module_type,
|
||||
motion_module_kwargs=motion_module_kwargs,
|
||||
)
|
||||
self.up_blocks.append(up_block)
|
||||
prev_output_channel = output_channel
|
||||
|
||||
# out
|
||||
if use_inflated_groupnorm:
|
||||
self.conv_norm_out = InflatedGroupNorm(
|
||||
num_channels=block_out_channels[0],
|
||||
num_groups=norm_num_groups,
|
||||
eps=norm_eps,
|
||||
)
|
||||
else:
|
||||
self.conv_norm_out = nn.GroupNorm(
|
||||
num_channels=block_out_channels[0],
|
||||
num_groups=norm_num_groups,
|
||||
eps=norm_eps,
|
||||
)
|
||||
self.conv_act = nn.SiLU()
|
||||
self.conv_out = InflatedConv3d(
|
||||
block_out_channels[0], out_channels, kernel_size=3, padding=1
|
||||
)
|
||||
|
||||
@property
|
||||
# Copied from diffusers.models.unet_2d_condition.UNet2DConditionModel.attn_processors
|
||||
def attn_processors(self) -> Dict[str, AttentionProcessor]:
|
||||
r"""
|
||||
Returns:
|
||||
`dict` of attention processors: A dictionary containing all attention processors used in the model with
|
||||
indexed by its weight name.
|
||||
"""
|
||||
# set recursively
|
||||
processors = {}
|
||||
|
||||
def fn_recursive_add_processors(
|
||||
name: str,
|
||||
module: torch.nn.Module,
|
||||
processors: Dict[str, AttentionProcessor],
|
||||
):
|
||||
if hasattr(module, "set_processor"):
|
||||
processors[f"{name}.processor"] = module.processor
|
||||
|
||||
for sub_name, child in module.named_children():
|
||||
if "temporal_transformer" not in sub_name:
|
||||
fn_recursive_add_processors(f"{name}.{sub_name}", child, processors)
|
||||
|
||||
return processors
|
||||
|
||||
for name, module in self.named_children():
|
||||
if "temporal_transformer" not in name:
|
||||
fn_recursive_add_processors(name, module, processors)
|
||||
|
||||
return processors
|
||||
|
||||
def set_attention_slice(self, slice_size):
|
||||
r"""
|
||||
Enable sliced attention computation.
|
||||
|
||||
When this option is enabled, the attention module will split the input tensor in slices, to compute attention
|
||||
in several steps. This is useful to save some memory in exchange for a small speed decrease.
|
||||
|
||||
Args:
|
||||
slice_size (`str` or `int` or `list(int)`, *optional*, defaults to `"auto"`):
|
||||
When `"auto"`, halves the input to the attention heads, so attention will be computed in two steps. If
|
||||
`"max"`, maxium amount of memory will be saved by running only one slice at a time. If a number is
|
||||
provided, uses as many slices as `attention_head_dim // slice_size`. In this case, `attention_head_dim`
|
||||
must be a multiple of `slice_size`.
|
||||
"""
|
||||
sliceable_head_dims = []
|
||||
|
||||
def fn_recursive_retrieve_slicable_dims(module: torch.nn.Module):
|
||||
if hasattr(module, "set_attention_slice"):
|
||||
sliceable_head_dims.append(module.sliceable_head_dim)
|
||||
|
||||
for child in module.children():
|
||||
fn_recursive_retrieve_slicable_dims(child)
|
||||
|
||||
# retrieve number of attention layers
|
||||
for module in self.children():
|
||||
fn_recursive_retrieve_slicable_dims(module)
|
||||
|
||||
num_slicable_layers = len(sliceable_head_dims)
|
||||
|
||||
if slice_size == "auto":
|
||||
# half the attention head size is usually a good trade-off between
|
||||
# speed and memory
|
||||
slice_size = [dim // 2 for dim in sliceable_head_dims]
|
||||
elif slice_size == "max":
|
||||
# make smallest slice possible
|
||||
slice_size = num_slicable_layers * [1]
|
||||
|
||||
slice_size = (
|
||||
num_slicable_layers * [slice_size]
|
||||
if not isinstance(slice_size, list)
|
||||
else slice_size
|
||||
)
|
||||
|
||||
if len(slice_size) != len(sliceable_head_dims):
|
||||
raise ValueError(
|
||||
f"You have provided {len(slice_size)}, but {self.config} has {len(sliceable_head_dims)} different"
|
||||
f" attention layers. Make sure to match `len(slice_size)` to be {len(sliceable_head_dims)}."
|
||||
)
|
||||
|
||||
for i in range(len(slice_size)):
|
||||
size = slice_size[i]
|
||||
dim = sliceable_head_dims[i]
|
||||
if size is not None and size > dim:
|
||||
raise ValueError(f"size {size} has to be smaller or equal to {dim}.")
|
||||
|
||||
# Recursively walk through all the children.
|
||||
# Any children which exposes the set_attention_slice method
|
||||
# gets the message
|
||||
def fn_recursive_set_attention_slice(
|
||||
module: torch.nn.Module, slice_size: List[int]
|
||||
):
|
||||
if hasattr(module, "set_attention_slice"):
|
||||
module.set_attention_slice(slice_size.pop())
|
||||
|
||||
for child in module.children():
|
||||
fn_recursive_set_attention_slice(child, slice_size)
|
||||
|
||||
reversed_slice_size = list(reversed(slice_size))
|
||||
for module in self.children():
|
||||
fn_recursive_set_attention_slice(module, reversed_slice_size)
|
||||
|
||||
def _set_gradient_checkpointing(self, module, value=False):
|
||||
if hasattr(module, "gradient_checkpointing"):
|
||||
module.gradient_checkpointing = value
|
||||
|
||||
# Copied from diffusers.models.unet_2d_condition.UNet2DConditionModel.set_attn_processor
|
||||
def set_attn_processor(
|
||||
self, processor: Union[AttentionProcessor, Dict[str, AttentionProcessor]]
|
||||
):
|
||||
r"""
|
||||
Sets the attention processor to use to compute attention.
|
||||
|
||||
Parameters:
|
||||
processor (`dict` of `AttentionProcessor` or only `AttentionProcessor`):
|
||||
The instantiated processor class or a dictionary of processor classes that will be set as the processor
|
||||
for **all** `Attention` layers.
|
||||
|
||||
If `processor` is a dict, the key needs to define the path to the corresponding cross attention
|
||||
processor. This is strongly recommended when setting trainable attention processors.
|
||||
|
||||
"""
|
||||
count = len(self.attn_processors.keys())
|
||||
|
||||
if isinstance(processor, dict) and len(processor) != count:
|
||||
raise ValueError(
|
||||
f"A dict of processors was passed, but the number of processors {len(processor)} does not match the"
|
||||
f" number of attention layers: {count}. Please make sure to pass {count} processor classes."
|
||||
)
|
||||
|
||||
def fn_recursive_attn_processor(name: str, module: torch.nn.Module, processor):
|
||||
if hasattr(module, "set_processor"):
|
||||
if not isinstance(processor, dict):
|
||||
module.set_processor(processor)
|
||||
else:
|
||||
module.set_processor(processor.pop(f"{name}.processor"))
|
||||
|
||||
for sub_name, child in module.named_children():
|
||||
if "temporal_transformer" not in sub_name:
|
||||
fn_recursive_attn_processor(f"{name}.{sub_name}", child, processor)
|
||||
|
||||
for name, module in self.named_children():
|
||||
if "temporal_transformer" not in name:
|
||||
fn_recursive_attn_processor(name, module, processor)
|
||||
|
||||
def forward(
|
||||
self,
|
||||
sample: torch.FloatTensor,
|
||||
timestep: Union[torch.Tensor, float, int],
|
||||
encoder_hidden_states: torch.Tensor,
|
||||
class_labels: Optional[torch.Tensor] = None,
|
||||
pose_cond_fea: Optional[torch.Tensor] = None,
|
||||
attention_mask: Optional[torch.Tensor] = None,
|
||||
down_block_additional_residuals: Optional[Tuple[torch.Tensor]] = None,
|
||||
mid_block_additional_residual: Optional[torch.Tensor] = None,
|
||||
return_dict: bool = True,
|
||||
) -> Union[UNet3DConditionOutput, Tuple]:
|
||||
r"""
|
||||
Args:
|
||||
sample (`torch.FloatTensor`): (batch, channel, height, width) noisy inputs tensor
|
||||
timestep (`torch.FloatTensor` or `float` or `int`): (batch) timesteps
|
||||
encoder_hidden_states (`torch.FloatTensor`): (batch, sequence_length, feature_dim) encoder hidden states
|
||||
return_dict (`bool`, *optional*, defaults to `True`):
|
||||
Whether or not to return a [`models.unet_2d_condition.UNet2DConditionOutput`] instead of a plain tuple.
|
||||
|
||||
Returns:
|
||||
[`~models.unet_2d_condition.UNet2DConditionOutput`] or `tuple`:
|
||||
[`~models.unet_2d_condition.UNet2DConditionOutput`] if `return_dict` is True, otherwise a `tuple`. When
|
||||
returning a tuple, the first element is the sample tensor.
|
||||
"""
|
||||
# By default samples have to be AT least a multiple of the overall upsampling factor.
|
||||
# The overall upsampling factor is equal to 2 ** (# num of upsampling layears).
|
||||
# However, the upsampling interpolation output size can be forced to fit any upsampling size
|
||||
# on the fly if necessary.
|
||||
default_overall_up_factor = 2**self.num_upsamplers
|
||||
|
||||
# upsample size should be forwarded when sample is not a multiple of `default_overall_up_factor`
|
||||
forward_upsample_size = False
|
||||
upsample_size = None
|
||||
|
||||
if any(s % default_overall_up_factor != 0 for s in sample.shape[-2:]):
|
||||
logger.info("Forward upsample size to force interpolation output size.")
|
||||
forward_upsample_size = True
|
||||
|
||||
# prepare attention_mask
|
||||
if attention_mask is not None:
|
||||
attention_mask = (1 - attention_mask.to(sample.dtype)) * -10000.0
|
||||
attention_mask = attention_mask.unsqueeze(1)
|
||||
|
||||
# center input if necessary
|
||||
if self.config.center_input_sample:
|
||||
sample = 2 * sample - 1.0
|
||||
|
||||
# time
|
||||
timesteps = timestep
|
||||
if not torch.is_tensor(timesteps):
|
||||
# This would be a good case for the `match` statement (Python 3.10+)
|
||||
is_mps = sample.device.type == "mps"
|
||||
if isinstance(timestep, float):
|
||||
dtype = torch.float32 if is_mps else torch.float64
|
||||
else:
|
||||
dtype = torch.int32 if is_mps else torch.int64
|
||||
timesteps = torch.tensor([timesteps], dtype=dtype, device=sample.device)
|
||||
elif len(timesteps.shape) == 0:
|
||||
timesteps = timesteps[None].to(sample.device)
|
||||
|
||||
# broadcast to batch dimension in a way that's compatible with ONNX/Core ML
|
||||
timesteps = timesteps.expand(sample.shape[0])
|
||||
|
||||
t_emb = self.time_proj(timesteps)
|
||||
|
||||
# timesteps does not contain any weights and will always return f32 tensors
|
||||
# but time_embedding might actually be running in fp16. so we need to cast here.
|
||||
# there might be better ways to encapsulate this.
|
||||
t_emb = t_emb.to(dtype=self.dtype)
|
||||
emb = self.time_embedding(t_emb)
|
||||
|
||||
if self.class_embedding is not None:
|
||||
if class_labels is None:
|
||||
raise ValueError(
|
||||
"class_labels should be provided when num_class_embeds > 0"
|
||||
)
|
||||
|
||||
if self.config.class_embed_type == "timestep":
|
||||
class_labels = self.time_proj(class_labels)
|
||||
|
||||
class_emb = self.class_embedding(class_labels).to(dtype=self.dtype)
|
||||
emb = emb + class_emb
|
||||
|
||||
# pre-process
|
||||
sample = self.conv_in(sample)
|
||||
if pose_cond_fea is not None:
|
||||
sample = sample + pose_cond_fea
|
||||
|
||||
# down
|
||||
down_block_res_samples = (sample,)
|
||||
for downsample_block in self.down_blocks:
|
||||
if (
|
||||
hasattr(downsample_block, "has_cross_attention")
|
||||
and downsample_block.has_cross_attention
|
||||
):
|
||||
sample, res_samples = downsample_block(
|
||||
hidden_states=sample,
|
||||
temb=emb,
|
||||
encoder_hidden_states=encoder_hidden_states,
|
||||
attention_mask=attention_mask,
|
||||
)
|
||||
else:
|
||||
sample, res_samples = downsample_block(
|
||||
hidden_states=sample,
|
||||
temb=emb,
|
||||
encoder_hidden_states=encoder_hidden_states,
|
||||
)
|
||||
|
||||
down_block_res_samples += res_samples
|
||||
|
||||
if down_block_additional_residuals is not None:
|
||||
new_down_block_res_samples = ()
|
||||
|
||||
for down_block_res_sample, down_block_additional_residual in zip(
|
||||
down_block_res_samples, down_block_additional_residuals
|
||||
):
|
||||
down_block_res_sample = (
|
||||
down_block_res_sample + down_block_additional_residual
|
||||
)
|
||||
new_down_block_res_samples += (down_block_res_sample,)
|
||||
|
||||
down_block_res_samples = new_down_block_res_samples
|
||||
|
||||
# mid
|
||||
sample = self.mid_block(
|
||||
sample,
|
||||
emb,
|
||||
encoder_hidden_states=encoder_hidden_states,
|
||||
attention_mask=attention_mask,
|
||||
)
|
||||
|
||||
if mid_block_additional_residual is not None:
|
||||
sample = sample + mid_block_additional_residual
|
||||
|
||||
# up
|
||||
for i, upsample_block in enumerate(self.up_blocks):
|
||||
is_final_block = i == len(self.up_blocks) - 1
|
||||
|
||||
res_samples = down_block_res_samples[-len(upsample_block.resnets) :]
|
||||
down_block_res_samples = down_block_res_samples[
|
||||
: -len(upsample_block.resnets)
|
||||
]
|
||||
|
||||
# if we have not reached the final block and need to forward the
|
||||
# upsample size, we do it here
|
||||
if not is_final_block and forward_upsample_size:
|
||||
upsample_size = down_block_res_samples[-1].shape[2:]
|
||||
|
||||
if (
|
||||
hasattr(upsample_block, "has_cross_attention")
|
||||
and upsample_block.has_cross_attention
|
||||
):
|
||||
sample = upsample_block(
|
||||
hidden_states=sample,
|
||||
temb=emb,
|
||||
res_hidden_states_tuple=res_samples,
|
||||
encoder_hidden_states=encoder_hidden_states,
|
||||
upsample_size=upsample_size,
|
||||
attention_mask=attention_mask,
|
||||
)
|
||||
else:
|
||||
sample = upsample_block(
|
||||
hidden_states=sample,
|
||||
temb=emb,
|
||||
res_hidden_states_tuple=res_samples,
|
||||
upsample_size=upsample_size,
|
||||
encoder_hidden_states=encoder_hidden_states,
|
||||
)
|
||||
|
||||
# post-process
|
||||
sample = self.conv_norm_out(sample)
|
||||
sample = self.conv_act(sample)
|
||||
sample = self.conv_out(sample)
|
||||
|
||||
if not return_dict:
|
||||
return (sample,)
|
||||
|
||||
return UNet3DConditionOutput(sample=sample)
|
||||
|
||||
@classmethod
|
||||
def from_pretrained_2d(
|
||||
cls,
|
||||
pretrained_model_path: PathLike,
|
||||
motion_module_path: PathLike,
|
||||
subfolder=None,
|
||||
unet_additional_kwargs=None,
|
||||
mm_zero_proj_out=False,
|
||||
):
|
||||
pretrained_model_path = Path(pretrained_model_path)
|
||||
motion_module_path = Path(motion_module_path)
|
||||
if subfolder is not None:
|
||||
pretrained_model_path = pretrained_model_path.joinpath(subfolder)
|
||||
logger.info(
|
||||
f"loaded temporal unet's pretrained weights from {pretrained_model_path} ..."
|
||||
)
|
||||
|
||||
config_file = pretrained_model_path / "config.json"
|
||||
if not (config_file.exists() and config_file.is_file()):
|
||||
raise RuntimeError(f"{config_file} does not exist or is not a file")
|
||||
|
||||
unet_config = cls.load_config(config_file)
|
||||
unet_config["_class_name"] = cls.__name__
|
||||
unet_config["down_block_types"] = [
|
||||
"CrossAttnDownBlock3D",
|
||||
"CrossAttnDownBlock3D",
|
||||
"CrossAttnDownBlock3D",
|
||||
"DownBlock3D",
|
||||
]
|
||||
unet_config["up_block_types"] = [
|
||||
"UpBlock3D",
|
||||
"CrossAttnUpBlock3D",
|
||||
"CrossAttnUpBlock3D",
|
||||
"CrossAttnUpBlock3D",
|
||||
]
|
||||
unet_config["mid_block_type"] = "UNetMidBlock3DCrossAttn"
|
||||
|
||||
model = cls.from_config(unet_config, **unet_additional_kwargs)
|
||||
# load the vanilla weights
|
||||
if pretrained_model_path.joinpath(SAFETENSORS_WEIGHTS_NAME).exists():
|
||||
logger.debug(
|
||||
f"loading safeTensors weights from {pretrained_model_path} ..."
|
||||
)
|
||||
state_dict = load_file(
|
||||
pretrained_model_path.joinpath(SAFETENSORS_WEIGHTS_NAME), device="cpu"
|
||||
)
|
||||
|
||||
elif pretrained_model_path.joinpath(WEIGHTS_NAME).exists():
|
||||
logger.debug(f"loading weights from {pretrained_model_path} ...")
|
||||
state_dict = torch.load(
|
||||
pretrained_model_path.joinpath(WEIGHTS_NAME),
|
||||
map_location="cpu",
|
||||
weights_only=True,
|
||||
)
|
||||
else:
|
||||
raise FileNotFoundError(f"no weights file found in {pretrained_model_path}")
|
||||
|
||||
# load the motion module weights
|
||||
if motion_module_path.exists() and motion_module_path.is_file():
|
||||
if motion_module_path.suffix.lower() in [".pth", ".pt", ".ckpt"]:
|
||||
logger.info(f"Load motion module params from {motion_module_path}")
|
||||
motion_state_dict = torch.load(
|
||||
motion_module_path, map_location="cpu", weights_only=True
|
||||
)
|
||||
elif motion_module_path.suffix.lower() == ".safetensors":
|
||||
motion_state_dict = load_file(motion_module_path, device="cpu")
|
||||
else:
|
||||
raise RuntimeError(
|
||||
f"unknown file format for motion module weights: {motion_module_path.suffix}"
|
||||
)
|
||||
if mm_zero_proj_out:
|
||||
logger.info(f"Zero initialize proj_out layers in motion module...")
|
||||
new_motion_state_dict = OrderedDict()
|
||||
for k in motion_state_dict:
|
||||
if "proj_out" in k:
|
||||
continue
|
||||
new_motion_state_dict[k] = motion_state_dict[k]
|
||||
motion_state_dict = new_motion_state_dict
|
||||
|
||||
# merge the state dicts
|
||||
state_dict.update(motion_state_dict)
|
||||
|
||||
# load the weights into the model
|
||||
m, u = model.load_state_dict(state_dict, strict=False)
|
||||
logger.debug(f"### missing keys: {len(m)}; \n### unexpected keys: {len(u)};")
|
||||
|
||||
params = [
|
||||
p.numel() if "temporal" in n else 0 for n, p in model.named_parameters()
|
||||
]
|
||||
logger.info(f"Loaded {sum(params) / 1e6}M-parameter motion module")
|
||||
|
||||
return model
|
||||
@@ -0,0 +1,872 @@
|
||||
# Adapted from https://github.com/huggingface/diffusers/blob/main/src/diffusers/models/unet_2d_blocks.py
|
||||
|
||||
import pdb
|
||||
|
||||
import torch
|
||||
from torch import nn
|
||||
|
||||
from .motion_module import get_motion_module
|
||||
|
||||
# from .motion_module import get_motion_module
|
||||
from .resnet import Downsample3D, ResnetBlock3D, Upsample3D
|
||||
from .transformer_3d import Transformer3DModel
|
||||
|
||||
|
||||
def get_down_block(
|
||||
down_block_type,
|
||||
num_layers,
|
||||
in_channels,
|
||||
out_channels,
|
||||
temb_channels,
|
||||
add_downsample,
|
||||
resnet_eps,
|
||||
resnet_act_fn,
|
||||
attn_num_head_channels,
|
||||
resnet_groups=None,
|
||||
cross_attention_dim=None,
|
||||
downsample_padding=None,
|
||||
dual_cross_attention=False,
|
||||
use_linear_projection=False,
|
||||
only_cross_attention=False,
|
||||
upcast_attention=False,
|
||||
resnet_time_scale_shift="default",
|
||||
unet_use_cross_frame_attention=None,
|
||||
unet_use_temporal_attention=None,
|
||||
use_inflated_groupnorm=None,
|
||||
use_motion_module=None,
|
||||
motion_module_type=None,
|
||||
motion_module_kwargs=None,
|
||||
):
|
||||
down_block_type = (
|
||||
down_block_type[7:]
|
||||
if down_block_type.startswith("UNetRes")
|
||||
else down_block_type
|
||||
)
|
||||
if down_block_type == "DownBlock3D":
|
||||
return DownBlock3D(
|
||||
num_layers=num_layers,
|
||||
in_channels=in_channels,
|
||||
out_channels=out_channels,
|
||||
temb_channels=temb_channels,
|
||||
add_downsample=add_downsample,
|
||||
resnet_eps=resnet_eps,
|
||||
resnet_act_fn=resnet_act_fn,
|
||||
resnet_groups=resnet_groups,
|
||||
downsample_padding=downsample_padding,
|
||||
resnet_time_scale_shift=resnet_time_scale_shift,
|
||||
use_inflated_groupnorm=use_inflated_groupnorm,
|
||||
use_motion_module=use_motion_module,
|
||||
motion_module_type=motion_module_type,
|
||||
motion_module_kwargs=motion_module_kwargs,
|
||||
)
|
||||
elif down_block_type == "CrossAttnDownBlock3D":
|
||||
if cross_attention_dim is None:
|
||||
raise ValueError(
|
||||
"cross_attention_dim must be specified for CrossAttnDownBlock3D"
|
||||
)
|
||||
return CrossAttnDownBlock3D(
|
||||
num_layers=num_layers,
|
||||
in_channels=in_channels,
|
||||
out_channels=out_channels,
|
||||
temb_channels=temb_channels,
|
||||
add_downsample=add_downsample,
|
||||
resnet_eps=resnet_eps,
|
||||
resnet_act_fn=resnet_act_fn,
|
||||
resnet_groups=resnet_groups,
|
||||
downsample_padding=downsample_padding,
|
||||
cross_attention_dim=cross_attention_dim,
|
||||
attn_num_head_channels=attn_num_head_channels,
|
||||
dual_cross_attention=dual_cross_attention,
|
||||
use_linear_projection=use_linear_projection,
|
||||
only_cross_attention=only_cross_attention,
|
||||
upcast_attention=upcast_attention,
|
||||
resnet_time_scale_shift=resnet_time_scale_shift,
|
||||
unet_use_cross_frame_attention=unet_use_cross_frame_attention,
|
||||
unet_use_temporal_attention=unet_use_temporal_attention,
|
||||
use_inflated_groupnorm=use_inflated_groupnorm,
|
||||
use_motion_module=use_motion_module,
|
||||
motion_module_type=motion_module_type,
|
||||
motion_module_kwargs=motion_module_kwargs,
|
||||
)
|
||||
raise ValueError(f"{down_block_type} does not exist.")
|
||||
|
||||
|
||||
def get_up_block(
|
||||
up_block_type,
|
||||
num_layers,
|
||||
in_channels,
|
||||
out_channels,
|
||||
prev_output_channel,
|
||||
temb_channels,
|
||||
add_upsample,
|
||||
resnet_eps,
|
||||
resnet_act_fn,
|
||||
attn_num_head_channels,
|
||||
resnet_groups=None,
|
||||
cross_attention_dim=None,
|
||||
dual_cross_attention=False,
|
||||
use_linear_projection=False,
|
||||
only_cross_attention=False,
|
||||
upcast_attention=False,
|
||||
resnet_time_scale_shift="default",
|
||||
unet_use_cross_frame_attention=None,
|
||||
unet_use_temporal_attention=None,
|
||||
use_inflated_groupnorm=None,
|
||||
use_motion_module=None,
|
||||
motion_module_type=None,
|
||||
motion_module_kwargs=None,
|
||||
):
|
||||
up_block_type = (
|
||||
up_block_type[7:] if up_block_type.startswith("UNetRes") else up_block_type
|
||||
)
|
||||
if up_block_type == "UpBlock3D":
|
||||
return UpBlock3D(
|
||||
num_layers=num_layers,
|
||||
in_channels=in_channels,
|
||||
out_channels=out_channels,
|
||||
prev_output_channel=prev_output_channel,
|
||||
temb_channels=temb_channels,
|
||||
add_upsample=add_upsample,
|
||||
resnet_eps=resnet_eps,
|
||||
resnet_act_fn=resnet_act_fn,
|
||||
resnet_groups=resnet_groups,
|
||||
resnet_time_scale_shift=resnet_time_scale_shift,
|
||||
use_inflated_groupnorm=use_inflated_groupnorm,
|
||||
use_motion_module=use_motion_module,
|
||||
motion_module_type=motion_module_type,
|
||||
motion_module_kwargs=motion_module_kwargs,
|
||||
)
|
||||
elif up_block_type == "CrossAttnUpBlock3D":
|
||||
if cross_attention_dim is None:
|
||||
raise ValueError(
|
||||
"cross_attention_dim must be specified for CrossAttnUpBlock3D"
|
||||
)
|
||||
return CrossAttnUpBlock3D(
|
||||
num_layers=num_layers,
|
||||
in_channels=in_channels,
|
||||
out_channels=out_channels,
|
||||
prev_output_channel=prev_output_channel,
|
||||
temb_channels=temb_channels,
|
||||
add_upsample=add_upsample,
|
||||
resnet_eps=resnet_eps,
|
||||
resnet_act_fn=resnet_act_fn,
|
||||
resnet_groups=resnet_groups,
|
||||
cross_attention_dim=cross_attention_dim,
|
||||
attn_num_head_channels=attn_num_head_channels,
|
||||
dual_cross_attention=dual_cross_attention,
|
||||
use_linear_projection=use_linear_projection,
|
||||
only_cross_attention=only_cross_attention,
|
||||
upcast_attention=upcast_attention,
|
||||
resnet_time_scale_shift=resnet_time_scale_shift,
|
||||
unet_use_cross_frame_attention=unet_use_cross_frame_attention,
|
||||
unet_use_temporal_attention=unet_use_temporal_attention,
|
||||
use_inflated_groupnorm=use_inflated_groupnorm,
|
||||
use_motion_module=use_motion_module,
|
||||
motion_module_type=motion_module_type,
|
||||
motion_module_kwargs=motion_module_kwargs,
|
||||
)
|
||||
raise ValueError(f"{up_block_type} does not exist.")
|
||||
|
||||
|
||||
class UNetMidBlock3DCrossAttn(nn.Module):
|
||||
def __init__(
|
||||
self,
|
||||
in_channels: int,
|
||||
temb_channels: int,
|
||||
dropout: float = 0.0,
|
||||
num_layers: int = 1,
|
||||
resnet_eps: float = 1e-6,
|
||||
resnet_time_scale_shift: str = "default",
|
||||
resnet_act_fn: str = "swish",
|
||||
resnet_groups: int = 32,
|
||||
resnet_pre_norm: bool = True,
|
||||
attn_num_head_channels=1,
|
||||
output_scale_factor=1.0,
|
||||
cross_attention_dim=1280,
|
||||
dual_cross_attention=False,
|
||||
use_linear_projection=False,
|
||||
upcast_attention=False,
|
||||
unet_use_cross_frame_attention=None,
|
||||
unet_use_temporal_attention=None,
|
||||
use_inflated_groupnorm=None,
|
||||
use_motion_module=None,
|
||||
motion_module_type=None,
|
||||
motion_module_kwargs=None,
|
||||
):
|
||||
super().__init__()
|
||||
|
||||
self.has_cross_attention = True
|
||||
self.attn_num_head_channels = attn_num_head_channels
|
||||
resnet_groups = (
|
||||
resnet_groups if resnet_groups is not None else min(in_channels // 4, 32)
|
||||
)
|
||||
|
||||
# there is always at least one resnet
|
||||
resnets = [
|
||||
ResnetBlock3D(
|
||||
in_channels=in_channels,
|
||||
out_channels=in_channels,
|
||||
temb_channels=temb_channels,
|
||||
eps=resnet_eps,
|
||||
groups=resnet_groups,
|
||||
dropout=dropout,
|
||||
time_embedding_norm=resnet_time_scale_shift,
|
||||
non_linearity=resnet_act_fn,
|
||||
output_scale_factor=output_scale_factor,
|
||||
pre_norm=resnet_pre_norm,
|
||||
use_inflated_groupnorm=use_inflated_groupnorm,
|
||||
)
|
||||
]
|
||||
attentions = []
|
||||
motion_modules = []
|
||||
|
||||
for _ in range(num_layers):
|
||||
if dual_cross_attention:
|
||||
raise NotImplementedError
|
||||
attentions.append(
|
||||
Transformer3DModel(
|
||||
attn_num_head_channels,
|
||||
in_channels // attn_num_head_channels,
|
||||
in_channels=in_channels,
|
||||
num_layers=1,
|
||||
cross_attention_dim=cross_attention_dim,
|
||||
norm_num_groups=resnet_groups,
|
||||
use_linear_projection=use_linear_projection,
|
||||
upcast_attention=upcast_attention,
|
||||
unet_use_cross_frame_attention=unet_use_cross_frame_attention,
|
||||
unet_use_temporal_attention=unet_use_temporal_attention,
|
||||
)
|
||||
)
|
||||
motion_modules.append(
|
||||
get_motion_module(
|
||||
in_channels=in_channels,
|
||||
motion_module_type=motion_module_type,
|
||||
motion_module_kwargs=motion_module_kwargs,
|
||||
)
|
||||
if use_motion_module
|
||||
else None
|
||||
)
|
||||
resnets.append(
|
||||
ResnetBlock3D(
|
||||
in_channels=in_channels,
|
||||
out_channels=in_channels,
|
||||
temb_channels=temb_channels,
|
||||
eps=resnet_eps,
|
||||
groups=resnet_groups,
|
||||
dropout=dropout,
|
||||
time_embedding_norm=resnet_time_scale_shift,
|
||||
non_linearity=resnet_act_fn,
|
||||
output_scale_factor=output_scale_factor,
|
||||
pre_norm=resnet_pre_norm,
|
||||
use_inflated_groupnorm=use_inflated_groupnorm,
|
||||
)
|
||||
)
|
||||
|
||||
self.attentions = nn.ModuleList(attentions)
|
||||
self.resnets = nn.ModuleList(resnets)
|
||||
self.motion_modules = nn.ModuleList(motion_modules)
|
||||
|
||||
def forward(
|
||||
self,
|
||||
hidden_states,
|
||||
temb=None,
|
||||
encoder_hidden_states=None,
|
||||
audio_cond_fea = None,
|
||||
attention_mask=None,
|
||||
):
|
||||
hidden_states = self.resnets[0](hidden_states, temb)
|
||||
for attn, resnet, motion_module in zip(
|
||||
self.attentions, self.resnets[1:], self.motion_modules
|
||||
):
|
||||
hidden_states = attn(
|
||||
hidden_states,
|
||||
encoder_hidden_states=encoder_hidden_states,
|
||||
audio_cond_fea = audio_cond_fea
|
||||
).sample
|
||||
hidden_states = (
|
||||
motion_module(
|
||||
hidden_states, temb, encoder_hidden_states=encoder_hidden_states
|
||||
)
|
||||
if motion_module is not None
|
||||
else hidden_states
|
||||
)
|
||||
hidden_states = resnet(hidden_states, temb)
|
||||
|
||||
return hidden_states
|
||||
|
||||
|
||||
class CrossAttnDownBlock3D(nn.Module):
|
||||
def __init__(
|
||||
self,
|
||||
in_channels: int,
|
||||
out_channels: int,
|
||||
temb_channels: int,
|
||||
dropout: float = 0.0,
|
||||
num_layers: int = 1,
|
||||
resnet_eps: float = 1e-6,
|
||||
resnet_time_scale_shift: str = "default",
|
||||
resnet_act_fn: str = "swish",
|
||||
resnet_groups: int = 32,
|
||||
resnet_pre_norm: bool = True,
|
||||
attn_num_head_channels=1,
|
||||
cross_attention_dim=1280,
|
||||
output_scale_factor=1.0,
|
||||
downsample_padding=1,
|
||||
add_downsample=True,
|
||||
dual_cross_attention=False,
|
||||
use_linear_projection=False,
|
||||
only_cross_attention=False,
|
||||
upcast_attention=False,
|
||||
unet_use_cross_frame_attention=None,
|
||||
unet_use_temporal_attention=None,
|
||||
use_inflated_groupnorm=None,
|
||||
use_motion_module=None,
|
||||
motion_module_type=None,
|
||||
motion_module_kwargs=None,
|
||||
):
|
||||
super().__init__()
|
||||
resnets = []
|
||||
attentions = []
|
||||
motion_modules = []
|
||||
|
||||
self.has_cross_attention = True
|
||||
self.attn_num_head_channels = attn_num_head_channels
|
||||
|
||||
for i in range(num_layers):
|
||||
in_channels = in_channels if i == 0 else out_channels
|
||||
resnets.append(
|
||||
ResnetBlock3D(
|
||||
in_channels=in_channels,
|
||||
out_channels=out_channels,
|
||||
temb_channels=temb_channels,
|
||||
eps=resnet_eps,
|
||||
groups=resnet_groups,
|
||||
dropout=dropout,
|
||||
time_embedding_norm=resnet_time_scale_shift,
|
||||
non_linearity=resnet_act_fn,
|
||||
output_scale_factor=output_scale_factor,
|
||||
pre_norm=resnet_pre_norm,
|
||||
use_inflated_groupnorm=use_inflated_groupnorm,
|
||||
)
|
||||
)
|
||||
if dual_cross_attention:
|
||||
raise NotImplementedError
|
||||
attentions.append(
|
||||
Transformer3DModel(
|
||||
attn_num_head_channels,
|
||||
out_channels // attn_num_head_channels,
|
||||
in_channels=out_channels,
|
||||
num_layers=1,
|
||||
cross_attention_dim=cross_attention_dim,
|
||||
norm_num_groups=resnet_groups,
|
||||
use_linear_projection=use_linear_projection,
|
||||
only_cross_attention=only_cross_attention,
|
||||
upcast_attention=upcast_attention,
|
||||
unet_use_cross_frame_attention=unet_use_cross_frame_attention,
|
||||
unet_use_temporal_attention=unet_use_temporal_attention,
|
||||
)
|
||||
)
|
||||
motion_modules.append(
|
||||
get_motion_module(
|
||||
in_channels=out_channels,
|
||||
motion_module_type=motion_module_type,
|
||||
motion_module_kwargs=motion_module_kwargs,
|
||||
)
|
||||
if use_motion_module
|
||||
else None
|
||||
)
|
||||
|
||||
self.attentions = nn.ModuleList(attentions)
|
||||
self.resnets = nn.ModuleList(resnets)
|
||||
self.motion_modules = nn.ModuleList(motion_modules)
|
||||
|
||||
if add_downsample:
|
||||
self.downsamplers = nn.ModuleList(
|
||||
[
|
||||
Downsample3D(
|
||||
out_channels,
|
||||
use_conv=True,
|
||||
out_channels=out_channels,
|
||||
padding=downsample_padding,
|
||||
name="op",
|
||||
)
|
||||
]
|
||||
)
|
||||
else:
|
||||
self.downsamplers = None
|
||||
|
||||
self.gradient_checkpointing = False
|
||||
|
||||
def forward(
|
||||
self,
|
||||
hidden_states,
|
||||
temb=None,
|
||||
encoder_hidden_states=None,
|
||||
audio_cond_fea=None,
|
||||
attention_mask=None,
|
||||
):
|
||||
output_states = ()
|
||||
|
||||
for i, (resnet, attn, motion_module) in enumerate(
|
||||
zip(self.resnets, self.attentions, self.motion_modules)
|
||||
):
|
||||
# self.gradient_checkpointing = False
|
||||
if self.training and self.gradient_checkpointing:
|
||||
|
||||
def create_custom_forward(module, return_dict=None):
|
||||
def custom_forward(*inputs):
|
||||
if return_dict is not None:
|
||||
return module(*inputs, return_dict=return_dict)
|
||||
else:
|
||||
return module(*inputs)
|
||||
|
||||
return custom_forward
|
||||
|
||||
hidden_states = torch.utils.checkpoint.checkpoint(
|
||||
create_custom_forward(resnet), hidden_states, temb
|
||||
)
|
||||
hidden_states = torch.utils.checkpoint.checkpoint(
|
||||
create_custom_forward(attn, return_dict=False),
|
||||
hidden_states,
|
||||
encoder_hidden_states,
|
||||
audio_cond_fea,
|
||||
)[0]
|
||||
|
||||
# add motion module
|
||||
hidden_states = (
|
||||
motion_module(
|
||||
hidden_states, temb, encoder_hidden_states=encoder_hidden_states
|
||||
)
|
||||
if motion_module is not None
|
||||
else hidden_states
|
||||
)
|
||||
|
||||
else:
|
||||
hidden_states = resnet(hidden_states, temb)
|
||||
hidden_states = attn(
|
||||
hidden_states,
|
||||
encoder_hidden_states=encoder_hidden_states,
|
||||
audio_cond_fea = audio_cond_fea,
|
||||
).sample
|
||||
|
||||
# add motion module
|
||||
hidden_states = (
|
||||
motion_module(
|
||||
hidden_states, temb, encoder_hidden_states=encoder_hidden_states
|
||||
)
|
||||
if motion_module is not None
|
||||
else hidden_states
|
||||
)
|
||||
|
||||
output_states += (hidden_states,)
|
||||
|
||||
if self.downsamplers is not None:
|
||||
for downsampler in self.downsamplers:
|
||||
hidden_states = downsampler(hidden_states)
|
||||
|
||||
output_states += (hidden_states,)
|
||||
|
||||
return hidden_states, output_states
|
||||
|
||||
|
||||
class DownBlock3D(nn.Module):
|
||||
def __init__(
|
||||
self,
|
||||
in_channels: int,
|
||||
out_channels: int,
|
||||
temb_channels: int,
|
||||
dropout: float = 0.0,
|
||||
num_layers: int = 1,
|
||||
resnet_eps: float = 1e-6,
|
||||
resnet_time_scale_shift: str = "default",
|
||||
resnet_act_fn: str = "swish",
|
||||
resnet_groups: int = 32,
|
||||
resnet_pre_norm: bool = True,
|
||||
output_scale_factor=1.0,
|
||||
add_downsample=True,
|
||||
downsample_padding=1,
|
||||
use_inflated_groupnorm=None,
|
||||
use_motion_module=None,
|
||||
motion_module_type=None,
|
||||
motion_module_kwargs=None,
|
||||
):
|
||||
super().__init__()
|
||||
resnets = []
|
||||
motion_modules = []
|
||||
|
||||
# use_motion_module = False
|
||||
for i in range(num_layers):
|
||||
in_channels = in_channels if i == 0 else out_channels
|
||||
resnets.append(
|
||||
ResnetBlock3D(
|
||||
in_channels=in_channels,
|
||||
out_channels=out_channels,
|
||||
temb_channels=temb_channels,
|
||||
eps=resnet_eps,
|
||||
groups=resnet_groups,
|
||||
dropout=dropout,
|
||||
time_embedding_norm=resnet_time_scale_shift,
|
||||
non_linearity=resnet_act_fn,
|
||||
output_scale_factor=output_scale_factor,
|
||||
pre_norm=resnet_pre_norm,
|
||||
use_inflated_groupnorm=use_inflated_groupnorm,
|
||||
)
|
||||
)
|
||||
motion_modules.append(
|
||||
get_motion_module(
|
||||
in_channels=out_channels,
|
||||
motion_module_type=motion_module_type,
|
||||
motion_module_kwargs=motion_module_kwargs,
|
||||
)
|
||||
if use_motion_module
|
||||
else None
|
||||
)
|
||||
|
||||
self.resnets = nn.ModuleList(resnets)
|
||||
self.motion_modules = nn.ModuleList(motion_modules)
|
||||
|
||||
if add_downsample:
|
||||
self.downsamplers = nn.ModuleList(
|
||||
[
|
||||
Downsample3D(
|
||||
out_channels,
|
||||
use_conv=True,
|
||||
out_channels=out_channels,
|
||||
padding=downsample_padding,
|
||||
name="op",
|
||||
)
|
||||
]
|
||||
)
|
||||
else:
|
||||
self.downsamplers = None
|
||||
|
||||
self.gradient_checkpointing = False
|
||||
|
||||
def forward(self, hidden_states, temb=None, encoder_hidden_states=None):
|
||||
output_states = ()
|
||||
|
||||
for resnet, motion_module in zip(self.resnets, self.motion_modules):
|
||||
# print(f"DownBlock3D {self.gradient_checkpointing = }")
|
||||
if self.training and self.gradient_checkpointing:
|
||||
|
||||
def create_custom_forward(module):
|
||||
def custom_forward(*inputs):
|
||||
return module(*inputs)
|
||||
|
||||
return custom_forward
|
||||
|
||||
hidden_states = torch.utils.checkpoint.checkpoint(
|
||||
create_custom_forward(resnet), hidden_states, temb
|
||||
)
|
||||
if motion_module is not None:
|
||||
hidden_states = torch.utils.checkpoint.checkpoint(
|
||||
create_custom_forward(motion_module),
|
||||
hidden_states.requires_grad_(),
|
||||
temb,
|
||||
encoder_hidden_states,
|
||||
)
|
||||
else:
|
||||
hidden_states = resnet(hidden_states, temb)
|
||||
|
||||
# add motion module
|
||||
hidden_states = (
|
||||
motion_module(
|
||||
hidden_states, temb, encoder_hidden_states=encoder_hidden_states
|
||||
)
|
||||
if motion_module is not None
|
||||
else hidden_states
|
||||
)
|
||||
|
||||
output_states += (hidden_states,)
|
||||
|
||||
if self.downsamplers is not None:
|
||||
for downsampler in self.downsamplers:
|
||||
hidden_states = downsampler(hidden_states)
|
||||
|
||||
output_states += (hidden_states,)
|
||||
|
||||
return hidden_states, output_states
|
||||
|
||||
|
||||
class CrossAttnUpBlock3D(nn.Module):
|
||||
def __init__(
|
||||
self,
|
||||
in_channels: int,
|
||||
out_channels: int,
|
||||
prev_output_channel: int,
|
||||
temb_channels: int,
|
||||
dropout: float = 0.0,
|
||||
num_layers: int = 1,
|
||||
resnet_eps: float = 1e-6,
|
||||
resnet_time_scale_shift: str = "default",
|
||||
resnet_act_fn: str = "swish",
|
||||
resnet_groups: int = 32,
|
||||
resnet_pre_norm: bool = True,
|
||||
attn_num_head_channels=1,
|
||||
cross_attention_dim=1280,
|
||||
output_scale_factor=1.0,
|
||||
add_upsample=True,
|
||||
dual_cross_attention=False,
|
||||
use_linear_projection=False,
|
||||
only_cross_attention=False,
|
||||
upcast_attention=False,
|
||||
unet_use_cross_frame_attention=None,
|
||||
unet_use_temporal_attention=None,
|
||||
use_motion_module=None,
|
||||
use_inflated_groupnorm=None,
|
||||
motion_module_type=None,
|
||||
motion_module_kwargs=None,
|
||||
):
|
||||
super().__init__()
|
||||
resnets = []
|
||||
attentions = []
|
||||
motion_modules = []
|
||||
|
||||
self.has_cross_attention = True
|
||||
self.attn_num_head_channels = attn_num_head_channels
|
||||
|
||||
for i in range(num_layers):
|
||||
res_skip_channels = in_channels if (i == num_layers - 1) else out_channels
|
||||
resnet_in_channels = prev_output_channel if i == 0 else out_channels
|
||||
|
||||
resnets.append(
|
||||
ResnetBlock3D(
|
||||
in_channels=resnet_in_channels + res_skip_channels,
|
||||
out_channels=out_channels,
|
||||
temb_channels=temb_channels,
|
||||
eps=resnet_eps,
|
||||
groups=resnet_groups,
|
||||
dropout=dropout,
|
||||
time_embedding_norm=resnet_time_scale_shift,
|
||||
non_linearity=resnet_act_fn,
|
||||
output_scale_factor=output_scale_factor,
|
||||
pre_norm=resnet_pre_norm,
|
||||
use_inflated_groupnorm=use_inflated_groupnorm,
|
||||
)
|
||||
)
|
||||
if dual_cross_attention:
|
||||
raise NotImplementedError
|
||||
attentions.append(
|
||||
Transformer3DModel(
|
||||
attn_num_head_channels,
|
||||
out_channels // attn_num_head_channels,
|
||||
in_channels=out_channels,
|
||||
num_layers=1,
|
||||
cross_attention_dim=cross_attention_dim,
|
||||
norm_num_groups=resnet_groups,
|
||||
use_linear_projection=use_linear_projection,
|
||||
only_cross_attention=only_cross_attention,
|
||||
upcast_attention=upcast_attention,
|
||||
unet_use_cross_frame_attention=unet_use_cross_frame_attention,
|
||||
unet_use_temporal_attention=unet_use_temporal_attention,
|
||||
)
|
||||
)
|
||||
motion_modules.append(
|
||||
get_motion_module(
|
||||
in_channels=out_channels,
|
||||
motion_module_type=motion_module_type,
|
||||
motion_module_kwargs=motion_module_kwargs,
|
||||
)
|
||||
if use_motion_module
|
||||
else None
|
||||
)
|
||||
|
||||
self.attentions = nn.ModuleList(attentions)
|
||||
self.resnets = nn.ModuleList(resnets)
|
||||
self.motion_modules = nn.ModuleList(motion_modules)
|
||||
|
||||
if add_upsample:
|
||||
self.upsamplers = nn.ModuleList(
|
||||
[Upsample3D(out_channels, use_conv=True, out_channels=out_channels)]
|
||||
)
|
||||
else:
|
||||
self.upsamplers = None
|
||||
|
||||
self.gradient_checkpointing = False
|
||||
|
||||
def forward(
|
||||
self,
|
||||
hidden_states,
|
||||
res_hidden_states_tuple,
|
||||
temb=None,
|
||||
encoder_hidden_states=None,
|
||||
audio_cond_fea=None,
|
||||
upsample_size=None,
|
||||
attention_mask=None,
|
||||
):
|
||||
for i, (resnet, attn, motion_module) in enumerate(
|
||||
zip(self.resnets, self.attentions, self.motion_modules)
|
||||
):
|
||||
# pop res hidden states
|
||||
res_hidden_states = res_hidden_states_tuple[-1]
|
||||
res_hidden_states_tuple = res_hidden_states_tuple[:-1]
|
||||
hidden_states = torch.cat([hidden_states, res_hidden_states], dim=1)
|
||||
|
||||
if self.training and self.gradient_checkpointing:
|
||||
|
||||
def create_custom_forward(module, return_dict=None):
|
||||
def custom_forward(*inputs):
|
||||
if return_dict is not None:
|
||||
return module(*inputs, return_dict=return_dict)
|
||||
else:
|
||||
return module(*inputs)
|
||||
|
||||
return custom_forward
|
||||
|
||||
hidden_states = torch.utils.checkpoint.checkpoint(
|
||||
create_custom_forward(resnet), hidden_states, temb
|
||||
)
|
||||
hidden_states = attn(
|
||||
hidden_states,
|
||||
encoder_hidden_states=encoder_hidden_states,
|
||||
audio_cond_fea = audio_cond_fea
|
||||
).sample
|
||||
if motion_module is not None:
|
||||
|
||||
hidden_states = torch.utils.checkpoint.checkpoint(
|
||||
create_custom_forward(motion_module),
|
||||
hidden_states.requires_grad_(),
|
||||
temb,
|
||||
encoder_hidden_states,
|
||||
)
|
||||
else:
|
||||
hidden_states = resnet(hidden_states, temb)
|
||||
hidden_states = attn(
|
||||
hidden_states,
|
||||
encoder_hidden_states=encoder_hidden_states,
|
||||
audio_cond_fea = audio_cond_fea
|
||||
).sample
|
||||
# print("hidden_states shape1:", hidden_states.shape)
|
||||
# add motion module
|
||||
hidden_states = (
|
||||
motion_module(
|
||||
hidden_states, temb, encoder_hidden_states=encoder_hidden_states
|
||||
)
|
||||
if motion_module is not None
|
||||
else hidden_states
|
||||
)
|
||||
|
||||
if self.upsamplers is not None:
|
||||
for upsampler in self.upsamplers:
|
||||
hidden_states = upsampler(hidden_states, upsample_size)
|
||||
|
||||
return hidden_states
|
||||
|
||||
|
||||
class UpBlock3D(nn.Module):
|
||||
def __init__(
|
||||
self,
|
||||
in_channels: int,
|
||||
prev_output_channel: int,
|
||||
out_channels: int,
|
||||
temb_channels: int,
|
||||
dropout: float = 0.0,
|
||||
num_layers: int = 1,
|
||||
resnet_eps: float = 1e-6,
|
||||
resnet_time_scale_shift: str = "default",
|
||||
resnet_act_fn: str = "swish",
|
||||
resnet_groups: int = 32,
|
||||
resnet_pre_norm: bool = True,
|
||||
output_scale_factor=1.0,
|
||||
add_upsample=True,
|
||||
use_inflated_groupnorm=None,
|
||||
use_motion_module=None,
|
||||
motion_module_type=None,
|
||||
motion_module_kwargs=None,
|
||||
):
|
||||
super().__init__()
|
||||
resnets = []
|
||||
motion_modules = []
|
||||
|
||||
# use_motion_module = False
|
||||
for i in range(num_layers):
|
||||
res_skip_channels = in_channels if (i == num_layers - 1) else out_channels
|
||||
resnet_in_channels = prev_output_channel if i == 0 else out_channels
|
||||
|
||||
resnets.append(
|
||||
ResnetBlock3D(
|
||||
in_channels=resnet_in_channels + res_skip_channels,
|
||||
out_channels=out_channels,
|
||||
temb_channels=temb_channels,
|
||||
eps=resnet_eps,
|
||||
groups=resnet_groups,
|
||||
dropout=dropout,
|
||||
time_embedding_norm=resnet_time_scale_shift,
|
||||
non_linearity=resnet_act_fn,
|
||||
output_scale_factor=output_scale_factor,
|
||||
pre_norm=resnet_pre_norm,
|
||||
use_inflated_groupnorm=use_inflated_groupnorm,
|
||||
)
|
||||
)
|
||||
motion_modules.append(
|
||||
get_motion_module(
|
||||
in_channels=out_channels,
|
||||
motion_module_type=motion_module_type,
|
||||
motion_module_kwargs=motion_module_kwargs,
|
||||
)
|
||||
if use_motion_module
|
||||
else None
|
||||
)
|
||||
|
||||
self.resnets = nn.ModuleList(resnets)
|
||||
self.motion_modules = nn.ModuleList(motion_modules)
|
||||
|
||||
if add_upsample:
|
||||
self.upsamplers = nn.ModuleList(
|
||||
[Upsample3D(out_channels, use_conv=True, out_channels=out_channels)]
|
||||
)
|
||||
else:
|
||||
self.upsamplers = None
|
||||
|
||||
self.gradient_checkpointing = False
|
||||
|
||||
def forward(
|
||||
self,
|
||||
hidden_states,
|
||||
res_hidden_states_tuple,
|
||||
temb=None,
|
||||
upsample_size=None,
|
||||
encoder_hidden_states=None,
|
||||
):
|
||||
for resnet, motion_module in zip(self.resnets, self.motion_modules):
|
||||
# pop res hidden states
|
||||
res_hidden_states = res_hidden_states_tuple[-1]
|
||||
res_hidden_states_tuple = res_hidden_states_tuple[:-1]
|
||||
hidden_states = torch.cat([hidden_states, res_hidden_states], dim=1)
|
||||
|
||||
# print(f"UpBlock3D {self.gradient_checkpointing = }")
|
||||
if self.training and self.gradient_checkpointing:
|
||||
|
||||
def create_custom_forward(module):
|
||||
def custom_forward(*inputs):
|
||||
return module(*inputs)
|
||||
|
||||
return custom_forward
|
||||
|
||||
hidden_states = torch.utils.checkpoint.checkpoint(
|
||||
create_custom_forward(resnet), hidden_states, temb
|
||||
)
|
||||
if motion_module is not None:
|
||||
hidden_states = torch.utils.checkpoint.checkpoint(
|
||||
create_custom_forward(motion_module),
|
||||
hidden_states.requires_grad_(),
|
||||
temb,
|
||||
encoder_hidden_states,
|
||||
)
|
||||
else:
|
||||
# print("hidden_states shape1:", hidden_states.shape)
|
||||
hidden_states = resnet(hidden_states, temb)
|
||||
# print("hidden_states shape2:", hidden_states.shape)
|
||||
hidden_states = (
|
||||
motion_module(
|
||||
hidden_states, temb, encoder_hidden_states=encoder_hidden_states
|
||||
)
|
||||
if motion_module is not None
|
||||
else hidden_states
|
||||
)
|
||||
|
||||
if self.upsamplers is not None:
|
||||
for upsampler in self.upsamplers:
|
||||
hidden_states = upsampler(hidden_states, upsample_size)
|
||||
|
||||
return hidden_states
|
||||
@@ -0,0 +1,679 @@
|
||||
from collections import OrderedDict
|
||||
from dataclasses import dataclass
|
||||
from os import PathLike
|
||||
from pathlib import Path
|
||||
from typing import Dict, List, Optional, Tuple, Union
|
||||
|
||||
import torch
|
||||
import torch.nn as nn
|
||||
import torch.utils.checkpoint
|
||||
from diffusers.configuration_utils import ConfigMixin, register_to_config
|
||||
from diffusers.models.attention_processor import AttentionProcessor
|
||||
from diffusers.models.embeddings import TimestepEmbedding, Timesteps
|
||||
from diffusers.models.modeling_utils import ModelMixin
|
||||
from diffusers.utils import SAFETENSORS_WEIGHTS_NAME, WEIGHTS_NAME, BaseOutput, logging
|
||||
from safetensors.torch import load_file
|
||||
|
||||
from .resnet import InflatedConv3d, InflatedGroupNorm
|
||||
from .unet_3d_blocks import UNetMidBlock3DCrossAttn, get_down_block, get_up_block
|
||||
|
||||
logger = logging.get_logger(__name__) # pylint: disable=invalid-name
|
||||
|
||||
|
||||
@dataclass
|
||||
class UNet3DConditionOutput(BaseOutput):
|
||||
sample: torch.FloatTensor
|
||||
|
||||
|
||||
class EchoUNet3DConditionModel(ModelMixin, ConfigMixin):
|
||||
_supports_gradient_checkpointing = True
|
||||
|
||||
@register_to_config
|
||||
def __init__(
|
||||
self,
|
||||
sample_size: Optional[int] = None,
|
||||
in_channels: int = 4,
|
||||
out_channels: int = 4,
|
||||
center_input_sample: bool = False,
|
||||
flip_sin_to_cos: bool = True,
|
||||
freq_shift: int = 0,
|
||||
down_block_types: Tuple[str] = (
|
||||
"CrossAttnDownBlock3D",
|
||||
"CrossAttnDownBlock3D",
|
||||
"CrossAttnDownBlock3D",
|
||||
"DownBlock3D",
|
||||
),
|
||||
mid_block_type: str = "UNetMidBlock3DCrossAttn",
|
||||
up_block_types: Tuple[str] = (
|
||||
"UpBlock3D",
|
||||
"CrossAttnUpBlock3D",
|
||||
"CrossAttnUpBlock3D",
|
||||
"CrossAttnUpBlock3D",
|
||||
),
|
||||
only_cross_attention: Union[bool, Tuple[bool]] = False,
|
||||
block_out_channels: Tuple[int] = (320, 640, 1280, 1280),
|
||||
layers_per_block: int = 2,
|
||||
downsample_padding: int = 1,
|
||||
mid_block_scale_factor: float = 1,
|
||||
act_fn: str = "silu",
|
||||
norm_num_groups: int = 32,
|
||||
norm_eps: float = 1e-5,
|
||||
cross_attention_dim: int = 1280,
|
||||
attention_head_dim: Union[int, Tuple[int]] = 8,
|
||||
dual_cross_attention: bool = False,
|
||||
use_linear_projection: bool = False,
|
||||
class_embed_type: Optional[str] = None,
|
||||
num_class_embeds: Optional[int] = None,
|
||||
upcast_attention: bool = False,
|
||||
resnet_time_scale_shift: str = "default",
|
||||
use_inflated_groupnorm=False,
|
||||
# Additional
|
||||
use_motion_module=False,
|
||||
motion_module_resolutions=(1, 2, 4, 8),
|
||||
motion_module_mid_block=False,
|
||||
motion_module_decoder_only=False,
|
||||
motion_module_type=None,
|
||||
motion_module_kwargs={},
|
||||
unet_use_cross_frame_attention=None,
|
||||
unet_use_temporal_attention=None,
|
||||
):
|
||||
super().__init__()
|
||||
|
||||
self.sample_size = sample_size
|
||||
time_embed_dim = block_out_channels[0] * 4
|
||||
|
||||
# input
|
||||
self.conv_in = InflatedConv3d(
|
||||
in_channels, block_out_channels[0], kernel_size=3, padding=(1, 1)
|
||||
)
|
||||
|
||||
# time
|
||||
self.time_proj = Timesteps(block_out_channels[0], flip_sin_to_cos, freq_shift)
|
||||
timestep_input_dim = block_out_channels[0]
|
||||
|
||||
self.time_embedding = TimestepEmbedding(timestep_input_dim, time_embed_dim)
|
||||
|
||||
# class embedding
|
||||
if class_embed_type is None and num_class_embeds is not None:
|
||||
self.class_embedding = nn.Embedding(num_class_embeds, time_embed_dim)
|
||||
elif class_embed_type == "timestep":
|
||||
self.class_embedding = TimestepEmbedding(timestep_input_dim, time_embed_dim)
|
||||
elif class_embed_type == "identity":
|
||||
self.class_embedding = nn.Identity(time_embed_dim, time_embed_dim)
|
||||
else:
|
||||
self.class_embedding = None
|
||||
|
||||
self.down_blocks = nn.ModuleList([])
|
||||
self.mid_block = None
|
||||
self.up_blocks = nn.ModuleList([])
|
||||
|
||||
if isinstance(only_cross_attention, bool):
|
||||
only_cross_attention = [only_cross_attention] * len(down_block_types)
|
||||
|
||||
if isinstance(attention_head_dim, int):
|
||||
attention_head_dim = (attention_head_dim,) * len(down_block_types)
|
||||
|
||||
# down
|
||||
output_channel = block_out_channels[0]
|
||||
for i, down_block_type in enumerate(down_block_types):
|
||||
res = 2**i
|
||||
input_channel = output_channel
|
||||
output_channel = block_out_channels[i]
|
||||
is_final_block = i == len(block_out_channels) - 1
|
||||
|
||||
down_block = get_down_block(
|
||||
down_block_type,
|
||||
num_layers=layers_per_block,
|
||||
in_channels=input_channel,
|
||||
out_channels=output_channel,
|
||||
temb_channels=time_embed_dim,
|
||||
add_downsample=not is_final_block,
|
||||
resnet_eps=norm_eps,
|
||||
resnet_act_fn=act_fn,
|
||||
resnet_groups=norm_num_groups,
|
||||
cross_attention_dim=cross_attention_dim,
|
||||
attn_num_head_channels=attention_head_dim[i],
|
||||
downsample_padding=downsample_padding,
|
||||
dual_cross_attention=dual_cross_attention,
|
||||
use_linear_projection=use_linear_projection,
|
||||
only_cross_attention=only_cross_attention[i],
|
||||
upcast_attention=upcast_attention,
|
||||
resnet_time_scale_shift=resnet_time_scale_shift,
|
||||
unet_use_cross_frame_attention=unet_use_cross_frame_attention,
|
||||
unet_use_temporal_attention=unet_use_temporal_attention,
|
||||
use_inflated_groupnorm=use_inflated_groupnorm,
|
||||
use_motion_module=use_motion_module
|
||||
and (res in motion_module_resolutions)
|
||||
and (not motion_module_decoder_only),
|
||||
motion_module_type=motion_module_type,
|
||||
motion_module_kwargs=motion_module_kwargs,
|
||||
)
|
||||
self.down_blocks.append(down_block)
|
||||
|
||||
# mid
|
||||
if mid_block_type == "UNetMidBlock3DCrossAttn":
|
||||
self.mid_block = UNetMidBlock3DCrossAttn(
|
||||
in_channels=block_out_channels[-1],
|
||||
temb_channels=time_embed_dim,
|
||||
resnet_eps=norm_eps,
|
||||
resnet_act_fn=act_fn,
|
||||
output_scale_factor=mid_block_scale_factor,
|
||||
resnet_time_scale_shift=resnet_time_scale_shift,
|
||||
cross_attention_dim=cross_attention_dim,
|
||||
attn_num_head_channels=attention_head_dim[-1],
|
||||
resnet_groups=norm_num_groups,
|
||||
dual_cross_attention=dual_cross_attention,
|
||||
use_linear_projection=use_linear_projection,
|
||||
upcast_attention=upcast_attention,
|
||||
unet_use_cross_frame_attention=unet_use_cross_frame_attention,
|
||||
unet_use_temporal_attention=unet_use_temporal_attention,
|
||||
use_inflated_groupnorm=use_inflated_groupnorm,
|
||||
use_motion_module=use_motion_module and motion_module_mid_block,
|
||||
motion_module_type=motion_module_type,
|
||||
motion_module_kwargs=motion_module_kwargs,
|
||||
)
|
||||
else:
|
||||
raise ValueError(f"unknown mid_block_type : {mid_block_type}")
|
||||
self.unet_use_temporal_attention = unet_use_temporal_attention
|
||||
# count how many layers upsample the videos
|
||||
self.num_upsamplers = 0
|
||||
|
||||
# up
|
||||
reversed_block_out_channels = list(reversed(block_out_channels))
|
||||
reversed_attention_head_dim = list(reversed(attention_head_dim))
|
||||
only_cross_attention = list(reversed(only_cross_attention))
|
||||
output_channel = reversed_block_out_channels[0]
|
||||
for i, up_block_type in enumerate(up_block_types):
|
||||
res = 2 ** (3 - i)
|
||||
is_final_block = i == len(block_out_channels) - 1
|
||||
|
||||
prev_output_channel = output_channel
|
||||
output_channel = reversed_block_out_channels[i]
|
||||
input_channel = reversed_block_out_channels[
|
||||
min(i + 1, len(block_out_channels) - 1)
|
||||
]
|
||||
|
||||
# add upsample block for all BUT final layer
|
||||
if not is_final_block:
|
||||
add_upsample = True
|
||||
self.num_upsamplers += 1
|
||||
else:
|
||||
add_upsample = False
|
||||
|
||||
up_block = get_up_block(
|
||||
up_block_type,
|
||||
num_layers=layers_per_block + 1,
|
||||
in_channels=input_channel,
|
||||
out_channels=output_channel,
|
||||
prev_output_channel=prev_output_channel,
|
||||
temb_channels=time_embed_dim,
|
||||
add_upsample=add_upsample,
|
||||
resnet_eps=norm_eps,
|
||||
resnet_act_fn=act_fn,
|
||||
resnet_groups=norm_num_groups,
|
||||
cross_attention_dim=cross_attention_dim,
|
||||
attn_num_head_channels=reversed_attention_head_dim[i],
|
||||
dual_cross_attention=dual_cross_attention,
|
||||
use_linear_projection=use_linear_projection,
|
||||
only_cross_attention=only_cross_attention[i],
|
||||
upcast_attention=upcast_attention,
|
||||
resnet_time_scale_shift=resnet_time_scale_shift,
|
||||
unet_use_cross_frame_attention=unet_use_cross_frame_attention,
|
||||
unet_use_temporal_attention=unet_use_temporal_attention,
|
||||
use_inflated_groupnorm=use_inflated_groupnorm,
|
||||
use_motion_module=use_motion_module
|
||||
and (res in motion_module_resolutions),
|
||||
motion_module_type=motion_module_type,
|
||||
motion_module_kwargs=motion_module_kwargs,
|
||||
)
|
||||
self.up_blocks.append(up_block)
|
||||
prev_output_channel = output_channel
|
||||
|
||||
# out
|
||||
if use_inflated_groupnorm:
|
||||
self.conv_norm_out = InflatedGroupNorm(
|
||||
num_channels=block_out_channels[0],
|
||||
num_groups=norm_num_groups,
|
||||
eps=norm_eps,
|
||||
)
|
||||
else:
|
||||
self.conv_norm_out = nn.GroupNorm(
|
||||
num_channels=block_out_channels[0],
|
||||
num_groups=norm_num_groups,
|
||||
eps=norm_eps,
|
||||
)
|
||||
self.conv_act = nn.SiLU()
|
||||
self.conv_out = InflatedConv3d(
|
||||
block_out_channels[0], out_channels, kernel_size=3, padding=1
|
||||
)
|
||||
|
||||
@property
|
||||
# Copied from diffusers.models.unet_2d_condition.UNet2DConditionModel.attn_processors
|
||||
def attn_processors(self) -> Dict[str, AttentionProcessor]:
|
||||
r"""
|
||||
Returns:
|
||||
`dict` of attention processors: A dictionary containing all attention processors used in the model with
|
||||
indexed by its weight name.
|
||||
"""
|
||||
# set recursively
|
||||
processors = {}
|
||||
|
||||
def fn_recursive_add_processors(
|
||||
name: str,
|
||||
module: torch.nn.Module,
|
||||
processors: Dict[str, AttentionProcessor],
|
||||
):
|
||||
if hasattr(module, "set_processor"):
|
||||
processors[f"{name}.processor"] = module.processor
|
||||
|
||||
for sub_name, child in module.named_children():
|
||||
if "temporal_transformer" not in sub_name:
|
||||
fn_recursive_add_processors(f"{name}.{sub_name}", child, processors)
|
||||
|
||||
return processors
|
||||
|
||||
for name, module in self.named_children():
|
||||
if "temporal_transformer" not in name:
|
||||
fn_recursive_add_processors(name, module, processors)
|
||||
|
||||
return processors
|
||||
|
||||
def set_attention_slice(self, slice_size):
|
||||
r"""
|
||||
Enable sliced attention computation.
|
||||
|
||||
When this option is enabled, the attention module will split the input tensor in slices, to compute attention
|
||||
in several steps. This is useful to save some memory in exchange for a small speed decrease.
|
||||
|
||||
Args:
|
||||
slice_size (`str` or `int` or `list(int)`, *optional*, defaults to `"auto"`):
|
||||
When `"auto"`, halves the input to the attention heads, so attention will be computed in two steps. If
|
||||
`"max"`, maxium amount of memory will be saved by running only one slice at a time. If a number is
|
||||
provided, uses as many slices as `attention_head_dim // slice_size`. In this case, `attention_head_dim`
|
||||
must be a multiple of `slice_size`.
|
||||
"""
|
||||
sliceable_head_dims = []
|
||||
|
||||
def fn_recursive_retrieve_slicable_dims(module: torch.nn.Module):
|
||||
if hasattr(module, "set_attention_slice"):
|
||||
sliceable_head_dims.append(module.sliceable_head_dim)
|
||||
|
||||
for child in module.children():
|
||||
fn_recursive_retrieve_slicable_dims(child)
|
||||
|
||||
# retrieve number of attention layers
|
||||
for module in self.children():
|
||||
fn_recursive_retrieve_slicable_dims(module)
|
||||
|
||||
num_slicable_layers = len(sliceable_head_dims)
|
||||
|
||||
if slice_size == "auto":
|
||||
# half the attention head size is usually a good trade-off between
|
||||
# speed and memory
|
||||
slice_size = [dim // 2 for dim in sliceable_head_dims]
|
||||
elif slice_size == "max":
|
||||
# make smallest slice possible
|
||||
slice_size = num_slicable_layers * [1]
|
||||
|
||||
slice_size = (
|
||||
num_slicable_layers * [slice_size]
|
||||
if not isinstance(slice_size, list)
|
||||
else slice_size
|
||||
)
|
||||
|
||||
if len(slice_size) != len(sliceable_head_dims):
|
||||
raise ValueError(
|
||||
f"You have provided {len(slice_size)}, but {self.config} has {len(sliceable_head_dims)} different"
|
||||
f" attention layers. Make sure to match `len(slice_size)` to be {len(sliceable_head_dims)}."
|
||||
)
|
||||
|
||||
for i in range(len(slice_size)):
|
||||
size = slice_size[i]
|
||||
dim = sliceable_head_dims[i]
|
||||
if size is not None and size > dim:
|
||||
raise ValueError(f"size {size} has to be smaller or equal to {dim}.")
|
||||
|
||||
# Recursively walk through all the children.
|
||||
# Any children which exposes the set_attention_slice method
|
||||
# gets the message
|
||||
def fn_recursive_set_attention_slice(
|
||||
module: torch.nn.Module, slice_size: List[int]
|
||||
):
|
||||
if hasattr(module, "set_attention_slice"):
|
||||
module.set_attention_slice(slice_size.pop())
|
||||
|
||||
for child in module.children():
|
||||
fn_recursive_set_attention_slice(child, slice_size)
|
||||
|
||||
reversed_slice_size = list(reversed(slice_size))
|
||||
for module in self.children():
|
||||
fn_recursive_set_attention_slice(module, reversed_slice_size)
|
||||
|
||||
def _set_gradient_checkpointing(self, module, value=False):
|
||||
if hasattr(module, "gradient_checkpointing"):
|
||||
module.gradient_checkpointing = value
|
||||
|
||||
# Copied from diffusers.models.unet_2d_condition.UNet2DConditionModel.set_attn_processor
|
||||
def set_attn_processor(
|
||||
self, processor: Union[AttentionProcessor, Dict[str, AttentionProcessor]]
|
||||
):
|
||||
r"""
|
||||
Sets the attention processor to use to compute attention.
|
||||
|
||||
Parameters:
|
||||
processor (`dict` of `AttentionProcessor` or only `AttentionProcessor`):
|
||||
The instantiated processor class or a dictionary of processor classes that will be set as the processor
|
||||
for **all** `Attention` layers.
|
||||
|
||||
If `processor` is a dict, the key needs to define the path to the corresponding cross attention
|
||||
processor. This is strongly recommended when setting trainable attention processors.
|
||||
|
||||
"""
|
||||
count = len(self.attn_processors.keys())
|
||||
|
||||
if isinstance(processor, dict) and len(processor) != count:
|
||||
raise ValueError(
|
||||
f"A dict of processors was passed, but the number of processors {len(processor)} does not match the"
|
||||
f" number of attention layers: {count}. Please make sure to pass {count} processor classes."
|
||||
)
|
||||
|
||||
def fn_recursive_attn_processor(name: str, module: torch.nn.Module, processor):
|
||||
if hasattr(module, "set_processor"):
|
||||
if not isinstance(processor, dict):
|
||||
module.set_processor(processor)
|
||||
else:
|
||||
module.set_processor(processor.pop(f"{name}.processor"))
|
||||
|
||||
for sub_name, child in module.named_children():
|
||||
if "temporal_transformer" not in sub_name:
|
||||
fn_recursive_attn_processor(f"{name}.{sub_name}", child, processor)
|
||||
|
||||
for name, module in self.named_children():
|
||||
if "temporal_transformer" not in name:
|
||||
fn_recursive_attn_processor(name, module, processor)
|
||||
|
||||
def forward(
|
||||
self,
|
||||
sample: torch.FloatTensor,
|
||||
timestep: Union[torch.Tensor, float, int],
|
||||
encoder_hidden_states: torch.Tensor,
|
||||
class_labels: Optional[torch.Tensor] = None,
|
||||
audio_cond_fea: Optional[torch.Tensor] = None,
|
||||
attention_mask: Optional[torch.Tensor] = None,
|
||||
face_musk_fea: Optional[torch.Tensor] = None,
|
||||
down_block_additional_residuals: Optional[Tuple[torch.Tensor]] = None,
|
||||
mid_block_additional_residual: Optional[torch.Tensor] = None,
|
||||
return_dict: bool = True,
|
||||
) -> Union[UNet3DConditionOutput, Tuple]:
|
||||
r"""
|
||||
Args:
|
||||
sample (`torch.FloatTensor`): (batch, channel, height, width) noisy inputs tensor
|
||||
timestep (`torch.FloatTensor` or `float` or `int`): (batch) timesteps
|
||||
encoder_hidden_states (`torch.FloatTensor`): (batch, sequence_length, feature_dim) encoder hidden states
|
||||
return_dict (`bool`, *optional*, defaults to `True`):
|
||||
Whether or not to return a [`models.unet_2d_condition.UNet2DConditionOutput`] instead of a plain tuple.
|
||||
|
||||
Returns:
|
||||
[`~models.unet_2d_condition.UNet2DConditionOutput`] or `tuple`:
|
||||
[`~models.unet_2d_condition.UNet2DConditionOutput`] if `return_dict` is True, otherwise a `tuple`. When
|
||||
returning a tuple, the first element is the sample tensor.
|
||||
"""
|
||||
# By default samples have to be AT least a multiple of the overall upsampling factor.
|
||||
# The overall upsampling factor is equal to 2 ** (# num of upsampling layears).
|
||||
# However, the upsampling interpolation output size can be forced to fit any upsampling size
|
||||
# on the fly if necessary.
|
||||
default_overall_up_factor = 2**self.num_upsamplers
|
||||
|
||||
# upsample size should be forwarded when sample is not a multiple of `default_overall_up_factor`
|
||||
forward_upsample_size = False
|
||||
upsample_size = None
|
||||
|
||||
if any(s % default_overall_up_factor != 0 for s in sample.shape[-2:]):
|
||||
logger.info("Forward upsample size to force interpolation output size.")
|
||||
forward_upsample_size = True
|
||||
|
||||
# prepare attention_mask
|
||||
if attention_mask is not None:
|
||||
attention_mask = (1 - attention_mask.to(sample.dtype)) * -10000.0
|
||||
attention_mask = attention_mask.unsqueeze(1)
|
||||
|
||||
# center input if necessary
|
||||
if self.config.center_input_sample:
|
||||
sample = 2 * sample - 1.0
|
||||
|
||||
# time
|
||||
timesteps = timestep
|
||||
if not torch.is_tensor(timesteps):
|
||||
# This would be a good case for the `match` statement (Python 3.10+)
|
||||
is_mps = sample.device.type == "mps"
|
||||
if isinstance(timestep, float):
|
||||
dtype = torch.float32 if is_mps else torch.float64
|
||||
else:
|
||||
dtype = torch.int32 if is_mps else torch.int64
|
||||
timesteps = torch.tensor([timesteps], dtype=dtype, device=sample.device)
|
||||
elif len(timesteps.shape) == 0:
|
||||
timesteps = timesteps[None].to(sample.device)
|
||||
|
||||
# broadcast to batch dimension in a way that's compatible with ONNX/Core ML
|
||||
timesteps = timesteps.expand(sample.shape[0])
|
||||
|
||||
t_emb = self.time_proj(timesteps)
|
||||
|
||||
# timesteps does not contain any weights and will always return f32 tensors
|
||||
# but time_embedding might actually be running in fp16. so we need to cast here.
|
||||
# there might be better ways to encapsulate this.
|
||||
t_emb = t_emb.to(dtype=self.dtype)
|
||||
emb = self.time_embedding(t_emb)
|
||||
|
||||
if self.class_embedding is not None:
|
||||
if class_labels is None:
|
||||
raise ValueError(
|
||||
"class_labels should be provided when num_class_embeds > 0"
|
||||
)
|
||||
|
||||
if self.config.class_embed_type == "timestep":
|
||||
class_labels = self.time_proj(class_labels)
|
||||
|
||||
class_emb = self.class_embedding(class_labels).to(dtype=self.dtype)
|
||||
emb = emb + class_emb
|
||||
|
||||
# pre-process
|
||||
sample = self.conv_in(sample)
|
||||
if face_musk_fea is not None:
|
||||
# print(sample.shape, face_musk_fea.shape)
|
||||
sample = sample + face_musk_fea
|
||||
|
||||
# down
|
||||
down_block_res_samples = (sample,)
|
||||
for downsample_block in self.down_blocks:
|
||||
if (
|
||||
hasattr(downsample_block, "has_cross_attention")
|
||||
and downsample_block.has_cross_attention
|
||||
):
|
||||
sample, res_samples = downsample_block(
|
||||
hidden_states=sample,
|
||||
temb=emb,
|
||||
encoder_hidden_states=encoder_hidden_states,
|
||||
audio_cond_fea=audio_cond_fea,
|
||||
attention_mask=attention_mask,
|
||||
)
|
||||
else:
|
||||
sample, res_samples = downsample_block(
|
||||
hidden_states=sample,
|
||||
temb=emb,
|
||||
encoder_hidden_states=encoder_hidden_states,
|
||||
)
|
||||
|
||||
down_block_res_samples += res_samples
|
||||
|
||||
if down_block_additional_residuals is not None:
|
||||
new_down_block_res_samples = ()
|
||||
|
||||
for down_block_res_sample, down_block_additional_residual in zip(
|
||||
down_block_res_samples, down_block_additional_residuals
|
||||
):
|
||||
down_block_res_sample = (
|
||||
down_block_res_sample + down_block_additional_residual
|
||||
)
|
||||
new_down_block_res_samples += (down_block_res_sample,)
|
||||
|
||||
down_block_res_samples = new_down_block_res_samples
|
||||
|
||||
# mid
|
||||
sample = self.mid_block(
|
||||
sample,
|
||||
emb,
|
||||
encoder_hidden_states=encoder_hidden_states,
|
||||
audio_cond_fea=audio_cond_fea,
|
||||
attention_mask=attention_mask,
|
||||
)
|
||||
|
||||
if mid_block_additional_residual is not None:
|
||||
sample = sample + mid_block_additional_residual
|
||||
|
||||
# up
|
||||
for i, upsample_block in enumerate(self.up_blocks):
|
||||
is_final_block = i == len(self.up_blocks) - 1
|
||||
|
||||
res_samples = down_block_res_samples[-len(upsample_block.resnets) :]
|
||||
down_block_res_samples = down_block_res_samples[
|
||||
: -len(upsample_block.resnets)
|
||||
]
|
||||
|
||||
# if we have not reached the final block and need to forward the
|
||||
# upsample size, we do it here
|
||||
if not is_final_block and forward_upsample_size:
|
||||
upsample_size = down_block_res_samples[-1].shape[2:]
|
||||
|
||||
if (
|
||||
hasattr(upsample_block, "has_cross_attention")
|
||||
and upsample_block.has_cross_attention
|
||||
):
|
||||
sample = upsample_block(
|
||||
hidden_states=sample,
|
||||
temb=emb,
|
||||
res_hidden_states_tuple=res_samples,
|
||||
encoder_hidden_states=encoder_hidden_states,
|
||||
audio_cond_fea=audio_cond_fea,
|
||||
upsample_size=upsample_size,
|
||||
attention_mask=attention_mask,
|
||||
)
|
||||
else:
|
||||
sample = upsample_block(
|
||||
hidden_states=sample,
|
||||
temb=emb,
|
||||
res_hidden_states_tuple=res_samples,
|
||||
upsample_size=upsample_size,
|
||||
encoder_hidden_states=encoder_hidden_states,
|
||||
)
|
||||
|
||||
# post-process
|
||||
sample = self.conv_norm_out(sample)
|
||||
sample = self.conv_act(sample)
|
||||
sample = self.conv_out(sample)
|
||||
|
||||
if not return_dict:
|
||||
return (sample,)
|
||||
|
||||
return UNet3DConditionOutput(sample=sample)
|
||||
|
||||
@classmethod
|
||||
def from_pretrained_2d(
|
||||
cls,
|
||||
pretrained_model_path: PathLike,
|
||||
motion_module_path: PathLike,
|
||||
subfolder=None,
|
||||
unet_additional_kwargs=None,
|
||||
mm_zero_proj_out=False,
|
||||
):
|
||||
pretrained_model_path = Path(pretrained_model_path)
|
||||
motion_module_path = Path(motion_module_path)
|
||||
if subfolder is not None:
|
||||
pretrained_model_path = pretrained_model_path.joinpath(subfolder)
|
||||
logger.info(
|
||||
f"loaded temporal unet's pretrained weights from {pretrained_model_path} ..."
|
||||
)
|
||||
|
||||
config_file = pretrained_model_path / "config.json"
|
||||
if not (config_file.exists() and config_file.is_file()):
|
||||
raise RuntimeError(f"{config_file} does not exist or is not a file")
|
||||
|
||||
unet_config = cls.load_config(config_file)
|
||||
unet_config["_class_name"] = cls.__name__
|
||||
unet_config["down_block_types"] = [
|
||||
"CrossAttnDownBlock3D",
|
||||
"CrossAttnDownBlock3D",
|
||||
"CrossAttnDownBlock3D",
|
||||
"DownBlock3D",
|
||||
]
|
||||
unet_config["up_block_types"] = [
|
||||
"UpBlock3D",
|
||||
"CrossAttnUpBlock3D",
|
||||
"CrossAttnUpBlock3D",
|
||||
"CrossAttnUpBlock3D",
|
||||
]
|
||||
unet_config["mid_block_type"] = "UNetMidBlock3DCrossAttn"
|
||||
|
||||
model = cls.from_config(unet_config, **unet_additional_kwargs)
|
||||
# load the vanilla weights
|
||||
if pretrained_model_path.joinpath(SAFETENSORS_WEIGHTS_NAME).exists():
|
||||
logger.debug(
|
||||
f"loading safeTensors weights from {pretrained_model_path} ..."
|
||||
)
|
||||
state_dict = load_file(
|
||||
pretrained_model_path.joinpath(SAFETENSORS_WEIGHTS_NAME), device="cpu"
|
||||
)
|
||||
|
||||
elif pretrained_model_path.joinpath(WEIGHTS_NAME).exists():
|
||||
logger.debug(f"loading weights from {pretrained_model_path} ...")
|
||||
state_dict = torch.load(
|
||||
pretrained_model_path.joinpath(WEIGHTS_NAME),
|
||||
map_location="cpu",
|
||||
weights_only=True,
|
||||
)
|
||||
else:
|
||||
raise FileNotFoundError(f"no weights file found in {pretrained_model_path}")
|
||||
|
||||
# load the motion module weights
|
||||
if motion_module_path.exists() and motion_module_path.is_file():
|
||||
if motion_module_path.suffix.lower() in [".pth", ".pt", ".ckpt"]:
|
||||
logger.info(f"Load motion module params from {motion_module_path}")
|
||||
motion_state_dict = torch.load(
|
||||
motion_module_path, map_location="cpu", weights_only=True
|
||||
)
|
||||
elif motion_module_path.suffix.lower() == ".safetensors":
|
||||
motion_state_dict = load_file(motion_module_path, device="cpu")
|
||||
else:
|
||||
raise RuntimeError(
|
||||
f"unknown file format for motion module weights: {motion_module_path.suffix}"
|
||||
)
|
||||
if mm_zero_proj_out:
|
||||
logger.info(f"Zero initialize proj_out layers in motion module...")
|
||||
new_motion_state_dict = OrderedDict()
|
||||
for k in motion_state_dict:
|
||||
if "proj_out" in k:
|
||||
continue
|
||||
new_motion_state_dict[k] = motion_state_dict[k]
|
||||
motion_state_dict = new_motion_state_dict
|
||||
|
||||
# merge the state dicts
|
||||
state_dict.update(motion_state_dict)
|
||||
|
||||
init_state_dict = model.state_dict()
|
||||
|
||||
for key in init_state_dict:
|
||||
if state_dict.get(key) is not None:
|
||||
if state_dict.get(key).shape == init_state_dict[key].shape:
|
||||
init_state_dict[key] = state_dict.get(key)
|
||||
# else:
|
||||
# print(key, " tensor shape is not matched with origin SD.")
|
||||
# load the weights into the model
|
||||
m, u = model.load_state_dict(init_state_dict, strict=False)
|
||||
logger.debug(f"### missing keys: {len(m)}; \n### unexpected keys: {len(u)};")
|
||||
|
||||
params = [
|
||||
p.numel() if "temporal" in n else 0 for n, p in model.named_parameters()
|
||||
]
|
||||
logger.info(f"Loaded {sum(params) / 1e6}M-parameter motion module")
|
||||
|
||||
return model
|
||||
@@ -0,0 +1,131 @@
|
||||
import os
|
||||
from .whisper import load_model
|
||||
import numpy as np
|
||||
import time
|
||||
|
||||
class Audio2Feature():
|
||||
def __init__(self,
|
||||
whisper_model_type="tiny",
|
||||
model_path="./models/whisper/tiny.pt",
|
||||
device="cuda"):
|
||||
self.whisper_model_type = whisper_model_type
|
||||
self.model = load_model(model_path, device=device) #
|
||||
|
||||
def get_sliced_feature(self,
|
||||
feature_array,
|
||||
vid_idx,
|
||||
audio_feat_length=[2,2],
|
||||
fps=25):
|
||||
"""
|
||||
Get sliced features based on a given index
|
||||
:param feature_array:
|
||||
:param start_idx: the start index of the feature
|
||||
:param audio_feat_length:
|
||||
:return:
|
||||
"""
|
||||
length = len(feature_array)
|
||||
selected_feature = []
|
||||
selected_idx = []
|
||||
|
||||
center_idx = int(vid_idx*50/fps)
|
||||
left_idx = center_idx-audio_feat_length[0]*2
|
||||
right_idx = center_idx + (audio_feat_length[1]+1)*2
|
||||
|
||||
for idx in range(left_idx,right_idx):
|
||||
idx = max(0, idx)
|
||||
idx = min(length-1, idx)
|
||||
x = feature_array[idx]
|
||||
selected_feature.append(x)
|
||||
selected_idx.append(idx)
|
||||
|
||||
selected_feature = np.concatenate(selected_feature, axis=0)
|
||||
# print("before reshape:", selected_feature.shape)
|
||||
selected_feature = selected_feature.reshape(-1, 384)# 50*384
|
||||
return selected_feature,selected_idx
|
||||
|
||||
def get_sliced_feature_sparse(self,feature_array, vid_idx, audio_feat_length= [2,2],fps = 25):
|
||||
"""
|
||||
Get sliced features based on a given index
|
||||
:param feature_array:
|
||||
:param start_idx: the start index of the feature
|
||||
:param audio_feat_length:
|
||||
:return:
|
||||
"""
|
||||
length = len(feature_array)
|
||||
selected_feature = []
|
||||
selected_idx = []
|
||||
|
||||
for dt in range(-audio_feat_length[0],audio_feat_length[1]+1):
|
||||
left_idx = int((vid_idx+dt)*50/fps)
|
||||
if left_idx<1 or left_idx>length-1:
|
||||
left_idx = max(0, left_idx)
|
||||
left_idx = min(length-1, left_idx)
|
||||
|
||||
x = feature_array[left_idx]
|
||||
x = x[np.newaxis,:,:]
|
||||
x = np.repeat(x, 2, axis=0)
|
||||
selected_feature.append(x)
|
||||
selected_idx.append(left_idx)
|
||||
selected_idx.append(left_idx)
|
||||
else:
|
||||
x = feature_array[left_idx-1:left_idx+1]
|
||||
selected_feature.append(x)
|
||||
selected_idx.append(left_idx-1)
|
||||
selected_idx.append(left_idx)
|
||||
selected_feature = np.concatenate(selected_feature, axis=0)
|
||||
selected_feature = selected_feature.reshape(-1, 384)# 50*384
|
||||
return selected_feature,selected_idx
|
||||
|
||||
|
||||
def feature2chunks(self,feature_array,fps,audio_feat_length = [2,2]):
|
||||
whisper_chunks = []
|
||||
whisper_idx_multiplier = 50./fps
|
||||
i = 0
|
||||
print(f"video in {fps} FPS, audio idx in 50FPS")
|
||||
while 1:
|
||||
start_idx = int(i * whisper_idx_multiplier)
|
||||
selected_feature,selected_idx = self.get_sliced_feature(feature_array= feature_array,vid_idx = i,audio_feat_length=audio_feat_length,fps=fps)
|
||||
#print(f"i:{i},selected_idx {selected_idx}")
|
||||
whisper_chunks.append(selected_feature)
|
||||
i += 1
|
||||
if start_idx>len(feature_array):
|
||||
break
|
||||
|
||||
return np.array(whisper_chunks)
|
||||
|
||||
def audio2feat(self,audio_path):
|
||||
# get the sample rate of the audio
|
||||
result = self.model.transcribe(audio_path)
|
||||
embed_list = []
|
||||
for emb in result['segments']:
|
||||
encoder_embeddings = emb['encoder_embeddings']
|
||||
encoder_embeddings = encoder_embeddings.transpose(0,2,1,3)
|
||||
encoder_embeddings = encoder_embeddings.squeeze(0)
|
||||
start_idx = int(emb['start'])
|
||||
end_idx = int(emb['end'])
|
||||
emb_end_idx = int((end_idx - start_idx)/2)
|
||||
embed_list.append(encoder_embeddings[:emb_end_idx])
|
||||
concatenated_array = np.concatenate(embed_list, axis=0)
|
||||
return concatenated_array
|
||||
|
||||
def load_audio_model(model_path, device):
|
||||
audio_processor = Audio2Feature(model_path=model_path, device=device)
|
||||
return audio_processor
|
||||
|
||||
if __name__ == "__main__":
|
||||
audio_processor = Audio2Feature(model_path="../../models/whisper/whisper_tiny.pt")
|
||||
audio_path = "./test.mp3"
|
||||
array = audio_processor.audio2feat(audio_path)
|
||||
print(array.shape)
|
||||
fps = 25
|
||||
whisper_idx_multiplier = 50./fps
|
||||
|
||||
i = 0
|
||||
print(f"video in {fps} FPS, audio idx in 50FPS")
|
||||
while 1:
|
||||
start_idx = int(i * whisper_idx_multiplier)
|
||||
selected_feature,selected_idx = audio_processor.get_sliced_feature(feature_array= array,vid_idx = i,audio_feat_length=[2,2],fps=fps)
|
||||
print(f"video idx {i},\t audio idx {selected_idx},\t shape {selected_feature.shape}")
|
||||
i += 1
|
||||
if start_idx>len(array):
|
||||
break
|
||||
@@ -0,0 +1,116 @@
|
||||
import hashlib
|
||||
import io
|
||||
import os
|
||||
import urllib
|
||||
import warnings
|
||||
from typing import List, Optional, Union
|
||||
|
||||
import torch
|
||||
from tqdm import tqdm
|
||||
|
||||
from .audio import load_audio, log_mel_spectrogram, pad_or_trim
|
||||
from .decoding import DecodingOptions, DecodingResult, decode, detect_language
|
||||
from .model import Whisper, ModelDimensions
|
||||
from .transcribe import transcribe
|
||||
|
||||
|
||||
_MODELS = {
|
||||
"tiny.en": "https://openaipublic.azureedge.net/main/whisper/models/d3dd57d32accea0b295c96e26691aa14d8822fac7d9d27d5dc00b4ca2826dd03/tiny.en.pt",
|
||||
"tiny": "https://openaipublic.azureedge.net/main/whisper/models/65147644a518d12f04e32d6f3b26facc3f8dd46e5390956a9424a650c0ce22b9/tiny.pt",
|
||||
"base.en": "https://openaipublic.azureedge.net/main/whisper/models/25a8566e1d0c1e2231d1c762132cd20e0f96a85d16145c3a00adf5d1ac670ead/base.en.pt",
|
||||
"base": "https://openaipublic.azureedge.net/main/whisper/models/ed3a0b6b1c0edf879ad9b11b1af5a0e6ab5db9205f891f668f8b0e6c6326e34e/base.pt",
|
||||
"small.en": "https://openaipublic.azureedge.net/main/whisper/models/f953ad0fd29cacd07d5a9eda5624af0f6bcf2258be67c92b79389873d91e0872/small.en.pt",
|
||||
"small": "https://openaipublic.azureedge.net/main/whisper/models/9ecf779972d90ba49c06d968637d720dd632c55bbf19d441fb42bf17a411e794/small.pt",
|
||||
"medium.en": "https://openaipublic.azureedge.net/main/whisper/models/d7440d1dc186f76616474e0ff0b3b6b879abc9d1a4926b7adfa41db2d497ab4f/medium.en.pt",
|
||||
"medium": "https://openaipublic.azureedge.net/main/whisper/models/345ae4da62f9b3d59415adc60127b97c714f32e89e936602e85993674d08dcb1/medium.pt",
|
||||
"large": "https://openaipublic.azureedge.net/main/whisper/models/e4b87e7e0bf463eb8e6956e646f1e277e901512310def2c24bf0e11bd3c28e9a/large.pt",
|
||||
"large-v1": "https://openaipublic.azureedge.net/main/whisper/models/e4b87e7e0bf463eb8e6956e646f1e277e901512310def2c24bf0e11bd3c28e9a/large-v1.pt",
|
||||
"large-v2": "https://openaipublic.azureedge.net/main/whisper/models/81f7c96c852ee8fc832187b0132e569d6c3065a3252ed18e56effd0b6a73e524/large-v2.pt",
|
||||
"large-v3": "https://openaipublic.azureedge.net/main/whisper/models/e5b1a55b89c1367dacf97e3e19bfd829a01529dbfdeefa8caeb59b3f1b81dadb/large-v3.pt",
|
||||
}
|
||||
|
||||
|
||||
def _download(url: str, root: str, in_memory: bool) -> Union[bytes, str]:
|
||||
os.makedirs(root, exist_ok=True)
|
||||
|
||||
expected_sha256 = url.split("/")[-2]
|
||||
download_target = os.path.join(root, os.path.basename(url))
|
||||
|
||||
if os.path.exists(download_target) and not os.path.isfile(download_target):
|
||||
raise RuntimeError(f"{download_target} exists and is not a regular file")
|
||||
|
||||
if os.path.isfile(download_target):
|
||||
model_bytes = open(download_target, "rb").read()
|
||||
if hashlib.sha256(model_bytes).hexdigest() == expected_sha256:
|
||||
return model_bytes if in_memory else download_target
|
||||
else:
|
||||
warnings.warn(f"{download_target} exists, but the SHA256 checksum does not match; re-downloading the file")
|
||||
|
||||
with urllib.request.urlopen(url) as source, open(download_target, "wb") as output:
|
||||
with tqdm(total=int(source.info().get("Content-Length")), ncols=80, unit='iB', unit_scale=True, unit_divisor=1024) as loop:
|
||||
while True:
|
||||
buffer = source.read(8192)
|
||||
if not buffer:
|
||||
break
|
||||
|
||||
output.write(buffer)
|
||||
loop.update(len(buffer))
|
||||
|
||||
model_bytes = open(download_target, "rb").read()
|
||||
if hashlib.sha256(model_bytes).hexdigest() != expected_sha256:
|
||||
raise RuntimeError("Model has been downloaded but the SHA256 checksum does not not match. Please retry loading the model.")
|
||||
|
||||
return model_bytes if in_memory else download_target
|
||||
|
||||
|
||||
def available_models() -> List[str]:
|
||||
"""Returns the names of available models"""
|
||||
return list(_MODELS.keys())
|
||||
|
||||
|
||||
def load_model(name: str, device: Optional[Union[str, torch.device]] = None, download_root: str = None, in_memory: bool = False) -> Whisper:
|
||||
"""
|
||||
Load a Whisper ASR model
|
||||
|
||||
Parameters
|
||||
----------
|
||||
name : str
|
||||
one of the official model names listed by `whisper.available_models()`, or
|
||||
path to a model checkpoint containing the model dimensions and the model state_dict.
|
||||
device : Union[str, torch.device]
|
||||
the PyTorch device to put the model into
|
||||
download_root: str
|
||||
path to download the model files; by default, it uses "~/.cache/whisper"
|
||||
in_memory: bool
|
||||
whether to preload the model weights into host memory
|
||||
|
||||
Returns
|
||||
-------
|
||||
model : Whisper
|
||||
The Whisper ASR model instance
|
||||
"""
|
||||
|
||||
if device is None:
|
||||
device = "cuda" if torch.cuda.is_available() else "cpu"
|
||||
if download_root is None:
|
||||
download_root = os.getenv(
|
||||
"XDG_CACHE_HOME",
|
||||
os.path.join(os.path.expanduser("~"), ".cache", "whisper")
|
||||
)
|
||||
|
||||
if name in _MODELS:
|
||||
checkpoint_file = _download(_MODELS[name], download_root, in_memory)
|
||||
elif os.path.isfile(name):
|
||||
checkpoint_file = open(name, "rb").read() if in_memory else name
|
||||
else:
|
||||
raise RuntimeError(f"Model {name} not found; available models = {available_models()}")
|
||||
|
||||
with (io.BytesIO(checkpoint_file) if in_memory else open(checkpoint_file, "rb")) as fp:
|
||||
checkpoint = torch.load(fp, map_location=device)
|
||||
del checkpoint_file
|
||||
|
||||
dims = ModelDimensions(**checkpoint["dims"])
|
||||
model = Whisper(dims)
|
||||
model.load_state_dict(checkpoint["model_state_dict"])
|
||||
|
||||
return model.to(device)
|
||||
@@ -0,0 +1,4 @@
|
||||
from .transcribe import cli
|
||||
|
||||
|
||||
cli()
|
||||
File diff suppressed because it is too large
Load Diff
@@ -0,0 +1 @@
|
||||
{"bos_token": "<|endoftext|>", "eos_token": "<|endoftext|>", "unk_token": "<|endoftext|>"}
|
||||
@@ -0,0 +1 @@
|
||||
{"unk_token": "<|endoftext|>", "bos_token": "<|endoftext|>", "eos_token": "<|endoftext|>", "add_prefix_space": false, "model_max_length": 1024, "special_tokens_map_file": null, "name_or_path": "gpt2", "tokenizer_class": "GPT2Tokenizer"}
|
||||
File diff suppressed because one or more lines are too long
Binary file not shown.
@@ -0,0 +1 @@
|
||||
{"<|endoftext|>": 50257}
|
||||
File diff suppressed because it is too large
Load Diff
@@ -0,0 +1 @@
|
||||
{"bos_token": "<|endoftext|>", "eos_token": "<|endoftext|>", "unk_token": "<|endoftext|>"}
|
||||
@@ -0,0 +1 @@
|
||||
{"unk_token": {"content": "<|endoftext|>", "single_word": false, "lstrip": false, "rstrip": false, "normalized": true, "__type": "AddedToken"}, "bos_token": {"content": "<|endoftext|>", "single_word": false, "lstrip": false, "rstrip": false, "normalized": true, "__type": "AddedToken"}, "eos_token": {"content": "<|endoftext|>", "single_word": false, "lstrip": false, "rstrip": false, "normalized": true, "__type": "AddedToken"}, "add_prefix_space": false, "model_max_length": 1024, "special_tokens_map_file": null, "name_or_path": "multilingual", "errors": "replace", "tokenizer_class": "GPT2Tokenizer"}
|
||||
File diff suppressed because one or more lines are too long
@@ -0,0 +1,125 @@
|
||||
import os
|
||||
from functools import lru_cache
|
||||
from typing import Union
|
||||
|
||||
import ffmpeg
|
||||
import numpy as np
|
||||
import torch
|
||||
import torch.nn.functional as F
|
||||
|
||||
from .utils import exact_div
|
||||
|
||||
# hard-coded audio hyperparameters
|
||||
SAMPLE_RATE = 16000
|
||||
N_FFT = 400
|
||||
N_MELS = 80
|
||||
HOP_LENGTH = 160
|
||||
CHUNK_LENGTH = 30
|
||||
N_SAMPLES = CHUNK_LENGTH * SAMPLE_RATE # 480000: number of samples in a chunk
|
||||
N_FRAMES = exact_div(N_SAMPLES, HOP_LENGTH) # 3000: number of frames in a mel spectrogram input
|
||||
|
||||
|
||||
def load_audio(file: str, sr: int = SAMPLE_RATE):
|
||||
"""
|
||||
Open an audio file and read as mono waveform, resampling as necessary
|
||||
|
||||
Parameters
|
||||
----------
|
||||
file: str
|
||||
The audio file to open
|
||||
|
||||
sr: int
|
||||
The sample rate to resample the audio if necessary
|
||||
|
||||
Returns
|
||||
-------
|
||||
A NumPy array containing the audio waveform, in float32 dtype.
|
||||
"""
|
||||
try:
|
||||
# This launches a subprocess to decode audio while down-mixing and resampling as necessary.
|
||||
# Requires the ffmpeg CLI and `ffmpeg-python` package to be installed.
|
||||
out, _ = (
|
||||
ffmpeg.input(file, threads=0)
|
||||
.output("-", format="s16le", acodec="pcm_s16le", ac=1, ar=sr)
|
||||
.run(cmd=["ffmpeg", "-nostdin"], capture_stdout=True, capture_stderr=True)
|
||||
)
|
||||
except ffmpeg.Error as e:
|
||||
raise RuntimeError(f"Failed to load audio: {e.stderr.decode()}") from e
|
||||
|
||||
return np.frombuffer(out, np.int16).flatten().astype(np.float32) / 32768.0
|
||||
|
||||
|
||||
def pad_or_trim(array, length: int = N_SAMPLES, *, axis: int = -1):
|
||||
"""
|
||||
Pad or trim the audio array to N_SAMPLES, as expected by the encoder.
|
||||
"""
|
||||
if torch.is_tensor(array):
|
||||
if array.shape[axis] > length:
|
||||
array = array.index_select(dim=axis, index=torch.arange(length))
|
||||
|
||||
if array.shape[axis] < length:
|
||||
pad_widths = [(0, 0)] * array.ndim
|
||||
pad_widths[axis] = (0, length - array.shape[axis])
|
||||
array = F.pad(array, [pad for sizes in pad_widths[::-1] for pad in sizes])
|
||||
else:
|
||||
if array.shape[axis] > length:
|
||||
array = array.take(indices=range(length), axis=axis)
|
||||
|
||||
if array.shape[axis] < length:
|
||||
pad_widths = [(0, 0)] * array.ndim
|
||||
pad_widths[axis] = (0, length - array.shape[axis])
|
||||
array = np.pad(array, pad_widths)
|
||||
|
||||
return array
|
||||
|
||||
|
||||
@lru_cache(maxsize=None)
|
||||
def mel_filters(device, n_mels: int = N_MELS) -> torch.Tensor:
|
||||
"""
|
||||
load the mel filterbank matrix for projecting STFT into a Mel spectrogram.
|
||||
Allows decoupling librosa dependency; saved using:
|
||||
|
||||
np.savez_compressed(
|
||||
"mel_filters.npz",
|
||||
mel_80=librosa.filters.mel(sr=16000, n_fft=400, n_mels=80),
|
||||
)
|
||||
"""
|
||||
assert n_mels == 80, f"Unsupported n_mels: {n_mels}"
|
||||
with np.load(os.path.join(os.path.dirname(__file__), "assets", "mel_filters.npz")) as f:
|
||||
return torch.from_numpy(f[f"mel_{n_mels}"]).to(device)
|
||||
|
||||
|
||||
def log_mel_spectrogram(audio: Union[str, np.ndarray, torch.Tensor], n_mels: int = N_MELS):
|
||||
"""
|
||||
Compute the log-Mel spectrogram of
|
||||
|
||||
Parameters
|
||||
----------
|
||||
audio: Union[str, np.ndarray, torch.Tensor], shape = (*)
|
||||
The path to audio or either a NumPy array or Tensor containing the audio waveform in 16 kHz
|
||||
|
||||
n_mels: int
|
||||
The number of Mel-frequency filters, only 80 is supported
|
||||
|
||||
Returns
|
||||
-------
|
||||
torch.Tensor, shape = (80, n_frames)
|
||||
A Tensor that contains the Mel spectrogram
|
||||
"""
|
||||
if not torch.is_tensor(audio):
|
||||
if isinstance(audio, str):
|
||||
audio = load_audio(audio)
|
||||
audio = torch.from_numpy(audio)
|
||||
|
||||
window = torch.hann_window(N_FFT).to(audio.device)
|
||||
stft = torch.stft(audio, N_FFT, HOP_LENGTH, window=window, return_complex=True)
|
||||
|
||||
magnitudes = stft[:, :-1].abs() ** 2
|
||||
|
||||
filters = mel_filters(audio.device, n_mels)
|
||||
mel_spec = filters @ magnitudes
|
||||
|
||||
log_spec = torch.clamp(mel_spec, min=1e-10).log10()
|
||||
log_spec = torch.maximum(log_spec, log_spec.max() - 8.0)
|
||||
log_spec = (log_spec + 4.0) / 4.0
|
||||
return log_spec
|
||||
@@ -0,0 +1,729 @@
|
||||
from dataclasses import dataclass, field
|
||||
from typing import Dict, List, Tuple, Iterable, Optional, Sequence, Union, TYPE_CHECKING
|
||||
|
||||
import numpy as np
|
||||
import torch
|
||||
import torch.nn.functional as F
|
||||
from torch import Tensor
|
||||
from torch.distributions import Categorical
|
||||
|
||||
from .audio import CHUNK_LENGTH
|
||||
from .tokenizer import Tokenizer, get_tokenizer
|
||||
from .utils import compression_ratio
|
||||
|
||||
if TYPE_CHECKING:
|
||||
from .model import Whisper
|
||||
|
||||
|
||||
@torch.no_grad()
|
||||
def detect_language(model: "Whisper", mel: Tensor, tokenizer: Tokenizer = None) -> Tuple[Tensor, List[dict]]:
|
||||
"""
|
||||
Detect the spoken language in the audio, and return them as list of strings, along with the ids
|
||||
of the most probable language tokens and the probability distribution over all language tokens.
|
||||
This is performed outside the main decode loop in order to not interfere with kv-caching.
|
||||
|
||||
Returns
|
||||
-------
|
||||
language_tokens : Tensor, shape = (n_audio,)
|
||||
ids of the most probable language tokens, which appears after the startoftranscript token.
|
||||
language_probs : List[Dict[str, float]], length = n_audio
|
||||
list of dictionaries containing the probability distribution over all languages.
|
||||
"""
|
||||
if tokenizer is None:
|
||||
tokenizer = get_tokenizer(model.is_multilingual)
|
||||
if tokenizer.language is None or tokenizer.language_token not in tokenizer.sot_sequence:
|
||||
raise ValueError(f"This model doesn't have language tokens so it can't perform lang id")
|
||||
|
||||
single = mel.ndim == 2
|
||||
if single:
|
||||
mel = mel.unsqueeze(0)
|
||||
|
||||
# skip encoder forward pass if already-encoded audio features were given
|
||||
if mel.shape[-2:] != (model.dims.n_audio_ctx, model.dims.n_audio_state):
|
||||
mel = model.encoder(mel)
|
||||
|
||||
# forward pass using a single token, startoftranscript
|
||||
n_audio = mel.shape[0]
|
||||
x = torch.tensor([[tokenizer.sot]] * n_audio).to(mel.device) # [n_audio, 1]
|
||||
logits = model.logits(x, mel)[:, 0]
|
||||
|
||||
# collect detected languages; suppress all non-language tokens
|
||||
mask = torch.ones(logits.shape[-1], dtype=torch.bool)
|
||||
mask[list(tokenizer.all_language_tokens)] = False
|
||||
logits[:, mask] = -np.inf
|
||||
language_tokens = logits.argmax(dim=-1)
|
||||
language_token_probs = logits.softmax(dim=-1).cpu()
|
||||
language_probs = [
|
||||
{
|
||||
c: language_token_probs[i, j].item()
|
||||
for j, c in zip(tokenizer.all_language_tokens, tokenizer.all_language_codes)
|
||||
}
|
||||
for i in range(n_audio)
|
||||
]
|
||||
|
||||
if single:
|
||||
language_tokens = language_tokens[0]
|
||||
language_probs = language_probs[0]
|
||||
|
||||
return language_tokens, language_probs
|
||||
|
||||
|
||||
@dataclass(frozen=True)
|
||||
class DecodingOptions:
|
||||
task: str = "transcribe" # whether to perform X->X "transcribe" or X->English "translate"
|
||||
language: Optional[str] = None # language that the audio is in; uses detected language if None
|
||||
|
||||
# sampling-related options
|
||||
temperature: float = 0.0
|
||||
sample_len: Optional[int] = None # maximum number of tokens to sample
|
||||
best_of: Optional[int] = None # number of independent samples to collect, when t > 0
|
||||
beam_size: Optional[int] = None # number of beams in beam search, when t == 0
|
||||
patience: Optional[float] = None # patience in beam search (https://arxiv.org/abs/2204.05424)
|
||||
|
||||
# options for ranking generations (either beams or best-of-N samples)
|
||||
length_penalty: Optional[float] = None # "alpha" in Google NMT, None defaults to length norm
|
||||
|
||||
# prompt, prefix, and token suppression
|
||||
prompt: Optional[Union[str, List[int]]] = None # text or tokens for the previous context
|
||||
prefix: Optional[Union[str, List[int]]] = None # text or tokens to prefix the current context
|
||||
suppress_blank: bool = True # this will suppress blank outputs
|
||||
|
||||
# list of tokens ids (or comma-separated token ids) to suppress
|
||||
# "-1" will suppress a set of symbols as defined in `tokenizer.non_speech_tokens()`
|
||||
suppress_tokens: Optional[Union[str, Iterable[int]]] = "-1"
|
||||
|
||||
# timestamp sampling options
|
||||
without_timestamps: bool = False # use <|notimestamps|> to sample text tokens only
|
||||
max_initial_timestamp: Optional[float] = 1.0 # the initial timestamp cannot be later than this
|
||||
|
||||
# implementation details
|
||||
fp16: bool = True # use fp16 for most of the calculation
|
||||
|
||||
|
||||
@dataclass(frozen=True)
|
||||
class DecodingResult:
|
||||
audio_features: Tensor
|
||||
language: str
|
||||
encoder_embeddings: np.ndarray
|
||||
decoder_embeddings: np.ndarray
|
||||
language_probs: Optional[Dict[str, float]] = None
|
||||
tokens: List[int] = field(default_factory=list)
|
||||
text: str = ""
|
||||
avg_logprob: float = np.nan
|
||||
no_speech_prob: float = np.nan
|
||||
temperature: float = np.nan
|
||||
compression_ratio: float = np.nan
|
||||
|
||||
|
||||
class Inference:
|
||||
def logits(self, tokens: Tensor, audio_features: Tensor) -> Tensor:
|
||||
"""Perform a forward pass on the decoder and return per-token logits"""
|
||||
raise NotImplementedError
|
||||
|
||||
def rearrange_kv_cache(self, source_indices) -> None:
|
||||
"""Update the key-value cache according to the updated beams"""
|
||||
raise NotImplementedError
|
||||
|
||||
def cleanup_caching(self) -> None:
|
||||
"""Clean up any resources or hooks after decoding is finished"""
|
||||
pass
|
||||
|
||||
|
||||
class PyTorchInference(Inference):
|
||||
def __init__(self, model: "Whisper", initial_token_length: int):
|
||||
self.model: "Whisper" = model
|
||||
self.initial_token_length = initial_token_length
|
||||
self.kv_cache = {}
|
||||
self.hooks = []
|
||||
|
||||
def logits(self, tokens: Tensor, audio_features: Tensor, include_embeddings=False) -> Tensor:
|
||||
if not self.kv_cache:
|
||||
self.kv_cache, self.hooks = self.model.install_kv_cache_hooks()
|
||||
|
||||
if tokens.shape[-1] > self.initial_token_length:
|
||||
# only need to use the last token except in the first forward pass
|
||||
tokens = tokens[:, -1:]
|
||||
|
||||
return_val = self.model.decoder(tokens, audio_features,
|
||||
kv_cache=self.kv_cache, include_embeddings=include_embeddings)
|
||||
return return_val
|
||||
|
||||
def cleanup_caching(self):
|
||||
for hook in self.hooks:
|
||||
hook.remove()
|
||||
|
||||
self.kv_cache = {}
|
||||
self.hooks = []
|
||||
|
||||
def rearrange_kv_cache(self, source_indices):
|
||||
for module, tensor in self.kv_cache.items():
|
||||
# update the key/value cache to contain the selected sequences
|
||||
self.kv_cache[module] = tensor[source_indices].detach()
|
||||
|
||||
|
||||
class SequenceRanker:
|
||||
def rank(self, tokens: List[List[Tensor]], sum_logprobs: List[List[float]]) -> List[int]:
|
||||
"""
|
||||
Given a list of groups of samples and their cumulative log probabilities,
|
||||
return the indices of the samples in each group to select as the final result
|
||||
"""
|
||||
raise NotImplementedError
|
||||
|
||||
|
||||
class MaximumLikelihoodRanker(SequenceRanker):
|
||||
"""
|
||||
Select the sample with the highest log probabilities, penalized using either
|
||||
a simple length normalization or Google NMT paper's length penalty
|
||||
"""
|
||||
|
||||
def __init__(self, length_penalty: Optional[float]):
|
||||
self.length_penalty = length_penalty
|
||||
|
||||
def rank(self, tokens: List[List[Tensor]], sum_logprobs: List[List[float]]):
|
||||
def scores(logprobs, lengths):
|
||||
result = []
|
||||
for logprob, length in zip(logprobs, lengths):
|
||||
if self.length_penalty is None:
|
||||
penalty = length
|
||||
else:
|
||||
# from the Google NMT paper
|
||||
penalty = ((5 + length) / 6) ** self.length_penalty
|
||||
result.append(logprob / penalty)
|
||||
return result
|
||||
|
||||
# get the sequence with the highest score
|
||||
lengths = [[len(t) for t in s] for s in tokens]
|
||||
return [np.argmax(scores(p, l)) for p, l in zip(sum_logprobs, lengths)]
|
||||
|
||||
|
||||
class TokenDecoder:
|
||||
def reset(self):
|
||||
"""Initialize any stateful variables for decoding a new sequence"""
|
||||
|
||||
def update(self, tokens: Tensor, logits: Tensor, sum_logprobs: Tensor) -> Tuple[Tensor, bool]:
|
||||
"""Specify how to select the next token, based on the current trace and logits
|
||||
|
||||
Parameters
|
||||
----------
|
||||
tokens : Tensor, shape = (n_batch, current_sequence_length)
|
||||
all tokens in the context so far, including the prefix and sot_sequence tokens
|
||||
|
||||
logits : Tensor, shape = (n_batch, vocab_size)
|
||||
per-token logits of the probability distribution at the current step
|
||||
|
||||
sum_logprobs : Tensor, shape = (n_batch)
|
||||
cumulative log probabilities for each sequence
|
||||
|
||||
Returns
|
||||
-------
|
||||
tokens : Tensor, shape = (n_batch, current_sequence_length + 1)
|
||||
the tokens, appended with the selected next token
|
||||
|
||||
completed : bool
|
||||
True if all sequences has reached the end of text
|
||||
|
||||
"""
|
||||
raise NotImplementedError
|
||||
|
||||
def finalize(
|
||||
self, tokens: Tensor, sum_logprobs: Tensor
|
||||
) -> Tuple[Sequence[Sequence[Tensor]], List[List[float]]]:
|
||||
"""Finalize search and return the final candidate sequences
|
||||
|
||||
Parameters
|
||||
----------
|
||||
tokens : Tensor, shape = (n_audio, n_group, current_sequence_length)
|
||||
all tokens in the context so far, including the prefix and sot_sequence
|
||||
|
||||
sum_logprobs : Tensor, shape = (n_audio, n_group)
|
||||
cumulative log probabilities for each sequence
|
||||
|
||||
Returns
|
||||
-------
|
||||
tokens : Sequence[Sequence[Tensor]], length = n_audio
|
||||
sequence of Tensors containing candidate token sequences, for each audio input
|
||||
|
||||
sum_logprobs : List[List[float]], length = n_audio
|
||||
sequence of cumulative log probabilities corresponding to the above
|
||||
|
||||
"""
|
||||
raise NotImplementedError
|
||||
|
||||
|
||||
class GreedyDecoder(TokenDecoder):
|
||||
def __init__(self, temperature: float, eot: int):
|
||||
self.temperature = temperature
|
||||
self.eot = eot
|
||||
|
||||
def update(self, tokens: Tensor, logits: Tensor, sum_logprobs: Tensor) -> Tuple[Tensor, bool]:
|
||||
temperature = self.temperature
|
||||
if temperature == 0:
|
||||
next_tokens = logits.argmax(dim=-1)
|
||||
else:
|
||||
next_tokens = Categorical(logits=logits / temperature).sample()
|
||||
|
||||
logprobs = F.log_softmax(logits.float(), dim=-1)
|
||||
current_logprobs = logprobs[torch.arange(logprobs.shape[0]), next_tokens]
|
||||
sum_logprobs += current_logprobs * (tokens[:, -1] != self.eot)
|
||||
|
||||
next_tokens[tokens[:, -1] == self.eot] = self.eot
|
||||
tokens = torch.cat([tokens, next_tokens[:, None]], dim=-1)
|
||||
|
||||
completed = (tokens[:, -1] == self.eot).all()
|
||||
return tokens, completed
|
||||
|
||||
def finalize(self, tokens: Tensor, sum_logprobs: Tensor):
|
||||
# make sure each sequence has at least one EOT token at the end
|
||||
tokens = F.pad(tokens, (0, 1), value=self.eot)
|
||||
return tokens, sum_logprobs.tolist()
|
||||
|
||||
|
||||
class BeamSearchDecoder(TokenDecoder):
|
||||
def __init__(self, beam_size: int, eot: int, inference: Inference, patience: Optional[float] = None):
|
||||
self.beam_size = beam_size
|
||||
self.eot = eot
|
||||
self.inference = inference
|
||||
self.patience = patience or 1.0
|
||||
self.max_candidates: int = round(beam_size * self.patience)
|
||||
self.finished_sequences = None
|
||||
|
||||
assert self.max_candidates > 0, f"Invalid beam size ({beam_size}) or patience ({patience})"
|
||||
|
||||
def reset(self):
|
||||
self.finished_sequences = None
|
||||
|
||||
def update(self, tokens: Tensor, logits: Tensor, sum_logprobs: Tensor) -> Tuple[Tensor, bool]:
|
||||
if tokens.shape[0] % self.beam_size != 0:
|
||||
raise ValueError(f"{tokens.shape}[0] % {self.beam_size} != 0")
|
||||
|
||||
n_audio = tokens.shape[0] // self.beam_size
|
||||
if self.finished_sequences is None: # for the first update
|
||||
self.finished_sequences = [{} for _ in range(n_audio)]
|
||||
|
||||
logprobs = F.log_softmax(logits.float(), dim=-1)
|
||||
next_tokens, source_indices, finished_sequences = [], [], []
|
||||
for i in range(n_audio):
|
||||
scores, sources, finished = {}, {}, {}
|
||||
|
||||
# STEP 1: calculate the cumulative log probabilities for possible candidates
|
||||
for j in range(self.beam_size):
|
||||
idx = i * self.beam_size + j
|
||||
prefix = tokens[idx].tolist()
|
||||
for logprob, token in zip(*logprobs[idx].topk(self.beam_size + 1)):
|
||||
new_logprob = (sum_logprobs[idx] + logprob).item()
|
||||
sequence = tuple(prefix + [token.item()])
|
||||
scores[sequence] = new_logprob
|
||||
sources[sequence] = idx
|
||||
|
||||
# STEP 2: rank the candidates and keep the top beam_size sequences for each audio
|
||||
saved = 0
|
||||
for sequence in sorted(scores, key=scores.get, reverse=True):
|
||||
if sequence[-1] == self.eot:
|
||||
finished[sequence] = scores[sequence]
|
||||
else:
|
||||
sum_logprobs[len(next_tokens)] = scores[sequence]
|
||||
next_tokens.append(sequence)
|
||||
source_indices.append(sources[sequence])
|
||||
|
||||
saved += 1
|
||||
if saved == self.beam_size:
|
||||
break
|
||||
|
||||
finished_sequences.append(finished)
|
||||
|
||||
tokens = torch.tensor(next_tokens, device=tokens.device)
|
||||
self.inference.rearrange_kv_cache(source_indices)
|
||||
|
||||
# add newly finished sequences to self.finished_sequences
|
||||
assert len(self.finished_sequences) == len(finished_sequences)
|
||||
for previously_finished, newly_finished in zip(self.finished_sequences, finished_sequences):
|
||||
for seq in sorted(newly_finished, key=newly_finished.get, reverse=True):
|
||||
if len(previously_finished) >= self.max_candidates:
|
||||
break # the candidate list is full
|
||||
previously_finished[seq] = newly_finished[seq]
|
||||
|
||||
# mark as completed if all audio has enough number of samples
|
||||
completed = all(
|
||||
len(sequences) >= self.max_candidates for sequences in self.finished_sequences
|
||||
)
|
||||
return tokens, completed
|
||||
|
||||
def finalize(self, preceding_tokens: Tensor, sum_logprobs: Tensor):
|
||||
# collect all finished sequences, including patience, and add unfinished ones if not enough
|
||||
sum_logprobs = sum_logprobs.cpu()
|
||||
for i, sequences in enumerate(self.finished_sequences):
|
||||
if len(sequences) < self.beam_size: # when not enough sequences are finished
|
||||
for j in list(np.argsort(sum_logprobs[i]))[::-1]:
|
||||
sequence = preceding_tokens[i, j].tolist() + [self.eot]
|
||||
sequences[tuple(sequence)] = sum_logprobs[i][j].item()
|
||||
if len(sequences) >= self.beam_size:
|
||||
break
|
||||
|
||||
tokens: List[List[Tensor]] = [
|
||||
[torch.tensor(seq) for seq in sequences.keys()] for sequences in self.finished_sequences
|
||||
]
|
||||
sum_logprobs: List[List[float]] = [
|
||||
list(sequences.values()) for sequences in self.finished_sequences
|
||||
]
|
||||
return tokens, sum_logprobs
|
||||
|
||||
|
||||
class LogitFilter:
|
||||
def apply(self, logits: Tensor, tokens: Tensor) -> None:
|
||||
"""Apply any filtering or masking to logits in-place
|
||||
|
||||
Parameters
|
||||
----------
|
||||
logits : Tensor, shape = (n_batch, vocab_size)
|
||||
per-token logits of the probability distribution at the current step
|
||||
|
||||
tokens : Tensor, shape = (n_batch, current_sequence_length)
|
||||
all tokens in the context so far, including the prefix and sot_sequence tokens
|
||||
|
||||
"""
|
||||
raise NotImplementedError
|
||||
|
||||
|
||||
class SuppressBlank(LogitFilter):
|
||||
def __init__(self, tokenizer: Tokenizer, sample_begin: int):
|
||||
self.tokenizer = tokenizer
|
||||
self.sample_begin = sample_begin
|
||||
|
||||
def apply(self, logits: Tensor, tokens: Tensor):
|
||||
if tokens.shape[1] == self.sample_begin:
|
||||
logits[:, self.tokenizer.encode(" ") + [self.tokenizer.eot]] = -np.inf
|
||||
|
||||
|
||||
class SuppressTokens(LogitFilter):
|
||||
def __init__(self, suppress_tokens: Sequence[int]):
|
||||
self.suppress_tokens = list(suppress_tokens)
|
||||
|
||||
def apply(self, logits: Tensor, tokens: Tensor):
|
||||
logits[:, self.suppress_tokens] = -np.inf
|
||||
|
||||
|
||||
class ApplyTimestampRules(LogitFilter):
|
||||
def __init__(
|
||||
self, tokenizer: Tokenizer, sample_begin: int, max_initial_timestamp_index: Optional[int]
|
||||
):
|
||||
self.tokenizer = tokenizer
|
||||
self.sample_begin = sample_begin
|
||||
self.max_initial_timestamp_index = max_initial_timestamp_index
|
||||
|
||||
def apply(self, logits: Tensor, tokens: Tensor):
|
||||
# suppress <|notimestamps|> which is handled by without_timestamps
|
||||
if self.tokenizer.no_timestamps is not None:
|
||||
logits[:, self.tokenizer.no_timestamps] = -np.inf
|
||||
|
||||
# timestamps have to appear in pairs, except directly before EOT; mask logits accordingly
|
||||
for k in range(tokens.shape[0]):
|
||||
seq = [t for t in tokens[k, self.sample_begin :].tolist()]
|
||||
last_was_timestamp = len(seq) >= 1 and seq[-1] >= self.tokenizer.timestamp_begin
|
||||
penultimate_was_timestamp = len(seq) < 2 or seq[-2] >= self.tokenizer.timestamp_begin
|
||||
|
||||
if last_was_timestamp:
|
||||
if penultimate_was_timestamp: # has to be non-timestamp
|
||||
logits[k, self.tokenizer.timestamp_begin :] = -np.inf
|
||||
else: # cannot be normal text tokens
|
||||
logits[k, : self.tokenizer.eot] = -np.inf
|
||||
|
||||
# apply the `max_initial_timestamp` option
|
||||
if tokens.shape[1] == self.sample_begin and self.max_initial_timestamp_index is not None:
|
||||
last_allowed = self.tokenizer.timestamp_begin + self.max_initial_timestamp_index
|
||||
logits[:, last_allowed + 1 :] = -np.inf
|
||||
|
||||
# if sum of probability over timestamps is above any other token, sample timestamp
|
||||
logprobs = F.log_softmax(logits.float(), dim=-1)
|
||||
for k in range(tokens.shape[0]):
|
||||
timestamp_logprob = logprobs[k, self.tokenizer.timestamp_begin :].logsumexp(dim=-1)
|
||||
max_text_token_logprob = logprobs[k, : self.tokenizer.timestamp_begin].max()
|
||||
if timestamp_logprob > max_text_token_logprob:
|
||||
logits[k, : self.tokenizer.timestamp_begin] = -np.inf
|
||||
|
||||
|
||||
class DecodingTask:
|
||||
inference: Inference
|
||||
sequence_ranker: SequenceRanker
|
||||
decoder: TokenDecoder
|
||||
logit_filters: List[LogitFilter]
|
||||
|
||||
def __init__(self, model: "Whisper", options: DecodingOptions):
|
||||
self.model = model
|
||||
|
||||
language = options.language or "en"
|
||||
tokenizer = get_tokenizer(model.is_multilingual, language=language, task=options.task)
|
||||
self.tokenizer: Tokenizer = tokenizer
|
||||
self.options: DecodingOptions = self._verify_options(options)
|
||||
|
||||
self.n_group: int = options.beam_size or options.best_of or 1
|
||||
self.n_ctx: int = model.dims.n_text_ctx
|
||||
self.sample_len: int = options.sample_len or model.dims.n_text_ctx // 2
|
||||
|
||||
self.sot_sequence: Tuple[int] = tokenizer.sot_sequence
|
||||
if self.options.without_timestamps:
|
||||
self.sot_sequence = tokenizer.sot_sequence_including_notimestamps
|
||||
|
||||
self.initial_tokens: Tuple[int] = self._get_initial_tokens()
|
||||
self.sample_begin: int = len(self.initial_tokens)
|
||||
self.sot_index: int = self.initial_tokens.index(tokenizer.sot)
|
||||
|
||||
# inference: implements the forward pass through the decoder, including kv caching
|
||||
self.inference = PyTorchInference(model, len(self.initial_tokens))
|
||||
|
||||
# sequence ranker: implements how to rank a group of sampled sequences
|
||||
self.sequence_ranker = MaximumLikelihoodRanker(options.length_penalty)
|
||||
|
||||
# decoder: implements how to select the next tokens, given the autoregressive distribution
|
||||
if options.beam_size is not None:
|
||||
self.decoder = BeamSearchDecoder(
|
||||
options.beam_size, tokenizer.eot, self.inference, options.patience
|
||||
)
|
||||
else:
|
||||
self.decoder = GreedyDecoder(options.temperature, tokenizer.eot)
|
||||
|
||||
# logit filters: applies various rules to suppress or penalize certain tokens
|
||||
self.logit_filters = []
|
||||
if self.options.suppress_blank:
|
||||
self.logit_filters.append(SuppressBlank(self.tokenizer, self.sample_begin))
|
||||
if self.options.suppress_tokens:
|
||||
self.logit_filters.append(SuppressTokens(self._get_suppress_tokens()))
|
||||
if not options.without_timestamps:
|
||||
precision = CHUNK_LENGTH / model.dims.n_audio_ctx # usually 0.02 seconds
|
||||
max_initial_timestamp_index = None
|
||||
if options.max_initial_timestamp:
|
||||
max_initial_timestamp_index = round(self.options.max_initial_timestamp / precision)
|
||||
self.logit_filters.append(
|
||||
ApplyTimestampRules(tokenizer, self.sample_begin, max_initial_timestamp_index)
|
||||
)
|
||||
|
||||
def _verify_options(self, options: DecodingOptions) -> DecodingOptions:
|
||||
if options.beam_size is not None and options.best_of is not None:
|
||||
raise ValueError("beam_size and best_of can't be given together")
|
||||
if options.temperature == 0:
|
||||
if options.best_of is not None:
|
||||
raise ValueError("best_of with greedy sampling (T=0) is not compatible")
|
||||
if options.patience is not None and options.beam_size is None:
|
||||
raise ValueError("patience requires beam_size to be given")
|
||||
if options.length_penalty is not None and not (0 <= options.length_penalty <= 1):
|
||||
raise ValueError("length_penalty (alpha) should be a value between 0 and 1")
|
||||
|
||||
return options
|
||||
|
||||
def _get_initial_tokens(self) -> Tuple[int]:
|
||||
tokens = list(self.sot_sequence)
|
||||
prefix = self.options.prefix
|
||||
prompt = self.options.prompt
|
||||
|
||||
if prefix:
|
||||
prefix_tokens = (
|
||||
self.tokenizer.encode(" " + prefix.strip()) if isinstance(prefix, str) else prefix
|
||||
)
|
||||
if self.sample_len is not None:
|
||||
max_prefix_len = self.n_ctx // 2 - self.sample_len
|
||||
prefix_tokens = prefix_tokens[-max_prefix_len:]
|
||||
tokens = tokens + prefix_tokens
|
||||
|
||||
if prompt:
|
||||
prompt_tokens = (
|
||||
self.tokenizer.encode(" " + prompt.strip()) if isinstance(prompt, str) else prompt
|
||||
)
|
||||
tokens = [self.tokenizer.sot_prev] + prompt_tokens[-(self.n_ctx // 2 - 1) :] + tokens
|
||||
|
||||
return tuple(tokens)
|
||||
|
||||
def _get_suppress_tokens(self) -> Tuple[int]:
|
||||
suppress_tokens = self.options.suppress_tokens
|
||||
|
||||
if isinstance(suppress_tokens, str):
|
||||
suppress_tokens = [int(t) for t in suppress_tokens.split(",")]
|
||||
|
||||
if -1 in suppress_tokens:
|
||||
suppress_tokens = [t for t in suppress_tokens if t >= 0]
|
||||
suppress_tokens.extend(self.tokenizer.non_speech_tokens)
|
||||
elif suppress_tokens is None or len(suppress_tokens) == 0:
|
||||
suppress_tokens = [] # interpret empty string as an empty list
|
||||
else:
|
||||
assert isinstance(suppress_tokens, list), "suppress_tokens must be a list"
|
||||
|
||||
suppress_tokens.extend(
|
||||
[self.tokenizer.sot, self.tokenizer.sot_prev, self.tokenizer.sot_lm]
|
||||
)
|
||||
if self.tokenizer.no_speech is not None:
|
||||
# no-speech probability is collected separately
|
||||
suppress_tokens.append(self.tokenizer.no_speech)
|
||||
|
||||
return tuple(sorted(set(suppress_tokens)))
|
||||
|
||||
def _get_audio_features(self, mel: Tensor, include_embeddings: bool = False):
|
||||
if self.options.fp16:
|
||||
mel = mel.half()
|
||||
|
||||
if mel.shape[-2:] == (self.model.dims.n_audio_ctx, self.model.dims.n_audio_state):
|
||||
# encoded audio features are given; skip audio encoding
|
||||
audio_features = mel
|
||||
else:
|
||||
result = self.model.encoder(mel, include_embeddings)
|
||||
if include_embeddings:
|
||||
audio_features, embeddings = result
|
||||
else:
|
||||
audio_features = result
|
||||
|
||||
if audio_features.dtype != (torch.float16 if self.options.fp16 else torch.float32):
|
||||
return TypeError(f"audio_features has an incorrect dtype: {audio_features.dtype}")
|
||||
|
||||
if include_embeddings:
|
||||
return audio_features, embeddings
|
||||
else:
|
||||
return audio_features
|
||||
|
||||
def _detect_language(self, audio_features: Tensor, tokens: Tensor):
|
||||
languages = [self.options.language] * audio_features.shape[0]
|
||||
lang_probs = None
|
||||
|
||||
if self.options.language is None or self.options.task == "lang_id":
|
||||
lang_tokens, lang_probs = self.model.detect_language(audio_features, self.tokenizer)
|
||||
languages = [max(probs, key=probs.get) for probs in lang_probs]
|
||||
if self.options.language is None:
|
||||
tokens[:, self.sot_index + 1] = lang_tokens # write language tokens
|
||||
|
||||
return languages, lang_probs
|
||||
|
||||
def _main_loop(self, audio_features: Tensor, tokens: Tensor):
|
||||
assert audio_features.shape[0] == tokens.shape[0]
|
||||
n_batch = tokens.shape[0]
|
||||
sum_logprobs: Tensor = torch.zeros(n_batch, device=audio_features.device)
|
||||
no_speech_probs = [np.nan] * n_batch
|
||||
|
||||
try:
|
||||
embeddings = []
|
||||
for i in range(self.sample_len):
|
||||
logits, token_embeddings = self.inference.logits(tokens, audio_features, include_embeddings=True)
|
||||
|
||||
if i == 0 and self.tokenizer.no_speech is not None: # save no_speech_probs
|
||||
probs_at_sot = logits[:, self.sot_index].float().softmax(dim=-1)
|
||||
no_speech_probs = probs_at_sot[:, self.tokenizer.no_speech].tolist()
|
||||
|
||||
# now we need to consider the logits at the last token only
|
||||
logits = logits[:, -1]
|
||||
token_embeddings = token_embeddings[:, :, -1]
|
||||
|
||||
# Append embeddings together
|
||||
embeddings.append(token_embeddings)
|
||||
|
||||
# apply the logit filters, e.g. for suppressing or applying penalty to
|
||||
for logit_filter in self.logit_filters:
|
||||
logit_filter.apply(logits, tokens)
|
||||
|
||||
# expand the tokens tensor with the selected next tokens
|
||||
tokens, completed = self.decoder.update(tokens, logits, sum_logprobs)
|
||||
|
||||
if completed or tokens.shape[-1] > self.n_ctx:
|
||||
break
|
||||
finally:
|
||||
if completed:
|
||||
embeddings = embeddings[:-1]
|
||||
embeddings = np.stack(embeddings, 2)
|
||||
self.inference.cleanup_caching()
|
||||
|
||||
return tokens, sum_logprobs, no_speech_probs, embeddings
|
||||
|
||||
@torch.no_grad()
|
||||
def run(self, mel: Tensor) -> List[DecodingResult]:
|
||||
self.decoder.reset()
|
||||
tokenizer: Tokenizer = self.tokenizer
|
||||
n_audio: int = mel.shape[0]
|
||||
|
||||
# encoder forward pass
|
||||
forward_pass: Tuple[Tensor, np.ndarray] = self._get_audio_features(mel, include_embeddings=True)
|
||||
audio_features, encoder_embeddings = forward_pass
|
||||
tokens: Tensor = torch.tensor([self.initial_tokens]).repeat(n_audio, 1)
|
||||
|
||||
# detect language if requested, overwriting the language token
|
||||
languages, language_probs = self._detect_language(audio_features, tokens)
|
||||
if self.options.task == "lang_id":
|
||||
return [
|
||||
DecodingResult(audio_features=features, language=language, language_probs=probs)
|
||||
for features, language, probs in zip(audio_features, languages, language_probs)
|
||||
]
|
||||
|
||||
# repeat the audio & text tensors by the group size, for beam search or best-of-n sampling
|
||||
audio_features = audio_features.repeat_interleave(self.n_group, dim=0)
|
||||
tokens = tokens.repeat_interleave(self.n_group, dim=0).to(audio_features.device)
|
||||
|
||||
# call the main sampling loop
|
||||
tokens, sum_logprobs, no_speech_probs, decoder_embeddings = self._main_loop(audio_features, tokens)
|
||||
|
||||
# reshape the tensors to have (n_audio, n_group) as the first two dimensions
|
||||
audio_features = audio_features[:: self.n_group]
|
||||
no_speech_probs = no_speech_probs[:: self.n_group]
|
||||
assert audio_features.shape[0] == len(no_speech_probs) == n_audio
|
||||
|
||||
tokens = tokens.reshape(n_audio, self.n_group, -1)
|
||||
sum_logprobs = sum_logprobs.reshape(n_audio, self.n_group)
|
||||
|
||||
# get the final candidates for each group, and slice between the first sampled token and EOT
|
||||
tokens, sum_logprobs = self.decoder.finalize(tokens, sum_logprobs)
|
||||
tokens: List[List[Tensor]] = [
|
||||
[t[self.sample_begin : (t == tokenizer.eot).nonzero()[0, 0]] for t in s] for s in tokens
|
||||
]
|
||||
|
||||
# select the top-ranked sample in each group
|
||||
selected = self.sequence_ranker.rank(tokens, sum_logprobs)
|
||||
tokens: List[List[int]] = [t[i].tolist() for i, t in zip(selected, tokens)]
|
||||
texts: List[str] = [tokenizer.decode(t).strip() for t in tokens]
|
||||
|
||||
sum_logprobs: List[float] = [lp[i] for i, lp in zip(selected, sum_logprobs)]
|
||||
avg_logprobs: List[float] = [lp / (len(t) + 1) for t, lp in zip(tokens, sum_logprobs)]
|
||||
|
||||
fields = (texts, languages, tokens, audio_features, avg_logprobs, no_speech_probs)
|
||||
if len(set(map(len, fields))) != 1:
|
||||
raise RuntimeError(f"inconsistent result lengths: {list(map(len, fields))}")
|
||||
|
||||
return [
|
||||
DecodingResult(
|
||||
audio_features=features,
|
||||
language=language,
|
||||
tokens=tokens,
|
||||
text=text,
|
||||
avg_logprob=avg_logprob,
|
||||
no_speech_prob=no_speech_prob,
|
||||
temperature=self.options.temperature,
|
||||
compression_ratio=compression_ratio(text),
|
||||
encoder_embeddings=encoder_embeddings,
|
||||
decoder_embeddings=decoder_embeddings
|
||||
)
|
||||
for text, language, tokens, features, avg_logprob, no_speech_prob in zip(*fields)
|
||||
]
|
||||
|
||||
|
||||
@torch.no_grad()
|
||||
def decode(model: "Whisper", mel: Tensor, options: DecodingOptions = DecodingOptions()) -> Union[DecodingResult, List[DecodingResult]]:
|
||||
"""
|
||||
Performs decoding of 30-second audio segment(s), provided as Mel spectrogram(s).
|
||||
|
||||
Parameters
|
||||
----------
|
||||
model: Whisper
|
||||
the Whisper model instance
|
||||
|
||||
mel: torch.Tensor, shape = (80, 3000) or (*, 80, 3000)
|
||||
A tensor containing the Mel spectrogram(s)
|
||||
|
||||
options: DecodingOptions
|
||||
A dataclass that contains all necessary options for decoding 30-second segments
|
||||
|
||||
Returns
|
||||
-------
|
||||
result: Union[DecodingResult, List[DecodingResult]]
|
||||
The result(s) of decoding contained in `DecodingResult` dataclass instance(s)
|
||||
"""
|
||||
single = mel.ndim == 2
|
||||
if single:
|
||||
mel = mel.unsqueeze(0)
|
||||
|
||||
result = DecodingTask(model, options).run(mel)
|
||||
|
||||
if single:
|
||||
result = result[0]
|
||||
|
||||
return result
|
||||
@@ -0,0 +1,290 @@
|
||||
from dataclasses import dataclass
|
||||
from typing import Dict
|
||||
from typing import Iterable, Optional
|
||||
|
||||
import numpy as np
|
||||
import torch
|
||||
import torch.nn.functional as F
|
||||
from torch import Tensor
|
||||
from torch import nn
|
||||
|
||||
from .transcribe import transcribe as transcribe_function
|
||||
from .decoding import detect_language as detect_language_function, decode as decode_function
|
||||
|
||||
|
||||
@dataclass
|
||||
class ModelDimensions:
|
||||
n_mels: int
|
||||
n_audio_ctx: int
|
||||
n_audio_state: int
|
||||
n_audio_head: int
|
||||
n_audio_layer: int
|
||||
n_vocab: int
|
||||
n_text_ctx: int
|
||||
n_text_state: int
|
||||
n_text_head: int
|
||||
n_text_layer: int
|
||||
|
||||
|
||||
class LayerNorm(nn.LayerNorm):
|
||||
def forward(self, x: Tensor) -> Tensor:
|
||||
return super().forward(x.float()).type(x.dtype)
|
||||
|
||||
|
||||
class Linear(nn.Linear):
|
||||
def forward(self, x: Tensor) -> Tensor:
|
||||
return F.linear(
|
||||
x, self.weight.to(x.dtype), None if self.bias is None else self.bias.to(x.dtype)
|
||||
)
|
||||
|
||||
|
||||
class Conv1d(nn.Conv1d):
|
||||
def _conv_forward(self, x: Tensor, weight: Tensor, bias: Optional[Tensor]) -> Tensor:
|
||||
return super()._conv_forward(
|
||||
x, weight.to(x.dtype), None if bias is None else bias.to(x.dtype)
|
||||
)
|
||||
|
||||
|
||||
def sinusoids(length, channels, max_timescale=10000):
|
||||
"""Returns sinusoids for positional embedding"""
|
||||
assert channels % 2 == 0
|
||||
log_timescale_increment = np.log(max_timescale) / (channels // 2 - 1)
|
||||
inv_timescales = torch.exp(-log_timescale_increment * torch.arange(channels // 2))
|
||||
scaled_time = torch.arange(length)[:, np.newaxis] * inv_timescales[np.newaxis, :]
|
||||
return torch.cat([torch.sin(scaled_time), torch.cos(scaled_time)], dim=1)
|
||||
|
||||
|
||||
class MultiHeadAttention(nn.Module):
|
||||
def __init__(self, n_state: int, n_head: int):
|
||||
super().__init__()
|
||||
self.n_head = n_head
|
||||
self.query = Linear(n_state, n_state)
|
||||
self.key = Linear(n_state, n_state, bias=False)
|
||||
self.value = Linear(n_state, n_state)
|
||||
self.out = Linear(n_state, n_state)
|
||||
|
||||
def forward(
|
||||
self,
|
||||
x: Tensor,
|
||||
xa: Optional[Tensor] = None,
|
||||
mask: Optional[Tensor] = None,
|
||||
kv_cache: Optional[dict] = None,
|
||||
):
|
||||
q = self.query(x)
|
||||
|
||||
if kv_cache is None or xa is None:
|
||||
# hooks, if installed (i.e. kv_cache is not None), will prepend the cached kv tensors;
|
||||
# otherwise, perform key/value projections for self- or cross-attention as usual.
|
||||
k = self.key(x if xa is None else xa)
|
||||
v = self.value(x if xa is None else xa)
|
||||
else:
|
||||
# for cross-attention, calculate keys and values once and reuse in subsequent calls.
|
||||
k = kv_cache.get(self.key, self.key(xa))
|
||||
v = kv_cache.get(self.value, self.value(xa))
|
||||
|
||||
wv = self.qkv_attention(q, k, v, mask)
|
||||
return self.out(wv)
|
||||
|
||||
def qkv_attention(self, q: Tensor, k: Tensor, v: Tensor, mask: Optional[Tensor] = None):
|
||||
n_batch, n_ctx, n_state = q.shape
|
||||
scale = (n_state // self.n_head) ** -0.25
|
||||
q = q.view(*q.shape[:2], self.n_head, -1).permute(0, 2, 1, 3) * scale
|
||||
k = k.view(*k.shape[:2], self.n_head, -1).permute(0, 2, 3, 1) * scale
|
||||
v = v.view(*v.shape[:2], self.n_head, -1).permute(0, 2, 1, 3)
|
||||
|
||||
qk = q @ k
|
||||
if mask is not None:
|
||||
qk = qk + mask[:n_ctx, :n_ctx]
|
||||
|
||||
w = F.softmax(qk.float(), dim=-1).to(q.dtype)
|
||||
return (w @ v).permute(0, 2, 1, 3).flatten(start_dim=2)
|
||||
|
||||
|
||||
class ResidualAttentionBlock(nn.Module):
|
||||
def __init__(self, n_state: int, n_head: int, cross_attention: bool = False):
|
||||
super().__init__()
|
||||
|
||||
self.attn = MultiHeadAttention(n_state, n_head)
|
||||
self.attn_ln = LayerNorm(n_state)
|
||||
|
||||
self.cross_attn = MultiHeadAttention(n_state, n_head) if cross_attention else None
|
||||
self.cross_attn_ln = LayerNorm(n_state) if cross_attention else None
|
||||
|
||||
n_mlp = n_state * 4
|
||||
self.mlp = nn.Sequential(Linear(n_state, n_mlp), nn.GELU(), Linear(n_mlp, n_state))
|
||||
self.mlp_ln = LayerNorm(n_state)
|
||||
|
||||
def forward(
|
||||
self,
|
||||
x: Tensor,
|
||||
xa: Optional[Tensor] = None,
|
||||
mask: Optional[Tensor] = None,
|
||||
kv_cache: Optional[dict] = None,
|
||||
):
|
||||
x = x + self.attn(self.attn_ln(x), mask=mask, kv_cache=kv_cache)
|
||||
if self.cross_attn:
|
||||
x = x + self.cross_attn(self.cross_attn_ln(x), xa, kv_cache=kv_cache)
|
||||
x = x + self.mlp(self.mlp_ln(x))
|
||||
return x
|
||||
|
||||
|
||||
class AudioEncoder(nn.Module):
|
||||
def __init__(self, n_mels: int, n_ctx: int, n_state: int, n_head: int, n_layer: int):
|
||||
super().__init__()
|
||||
self.conv1 = Conv1d(n_mels, n_state, kernel_size=3, padding=1)
|
||||
self.conv2 = Conv1d(n_state, n_state, kernel_size=3, stride=2, padding=1)
|
||||
self.register_buffer("positional_embedding", sinusoids(n_ctx, n_state))
|
||||
|
||||
self.blocks: Iterable[ResidualAttentionBlock] = nn.ModuleList(
|
||||
[ResidualAttentionBlock(n_state, n_head) for _ in range(n_layer)]
|
||||
)
|
||||
self.ln_post = LayerNorm(n_state)
|
||||
|
||||
def forward(self, x: Tensor, include_embeddings: bool = False):
|
||||
"""
|
||||
x : torch.Tensor, shape = (batch_size, n_mels, n_ctx)
|
||||
the mel spectrogram of the audio
|
||||
include_embeddings: bool
|
||||
whether to include intermediate steps in the output
|
||||
"""
|
||||
x = F.gelu(self.conv1(x))
|
||||
x = F.gelu(self.conv2(x))
|
||||
x = x.permute(0, 2, 1)
|
||||
|
||||
assert x.shape[1:] == self.positional_embedding.shape, "incorrect audio shape"
|
||||
x = (x + self.positional_embedding).to(x.dtype)
|
||||
|
||||
if include_embeddings:
|
||||
embeddings = [x.cpu().detach().numpy()]
|
||||
|
||||
for block in self.blocks:
|
||||
x = block(x)
|
||||
if include_embeddings:
|
||||
embeddings.append(x.cpu().detach().numpy())
|
||||
|
||||
x = self.ln_post(x)
|
||||
|
||||
if include_embeddings:
|
||||
embeddings = np.stack(embeddings, axis=1)
|
||||
return x, embeddings
|
||||
else:
|
||||
return x
|
||||
|
||||
|
||||
class TextDecoder(nn.Module):
|
||||
def __init__(self, n_vocab: int, n_ctx: int, n_state: int, n_head: int, n_layer: int):
|
||||
super().__init__()
|
||||
|
||||
self.token_embedding = nn.Embedding(n_vocab, n_state)
|
||||
self.positional_embedding = nn.Parameter(torch.empty(n_ctx, n_state))
|
||||
|
||||
self.blocks: Iterable[ResidualAttentionBlock] = nn.ModuleList(
|
||||
[ResidualAttentionBlock(n_state, n_head, cross_attention=True) for _ in range(n_layer)]
|
||||
)
|
||||
self.ln = LayerNorm(n_state)
|
||||
|
||||
mask = torch.empty(n_ctx, n_ctx).fill_(-np.inf).triu_(1)
|
||||
self.register_buffer("mask", mask, persistent=False)
|
||||
|
||||
def forward(self, x: Tensor, xa: Tensor, kv_cache: Optional[dict] = None, include_embeddings: bool = False):
|
||||
"""
|
||||
x : torch.LongTensor, shape = (batch_size, <= n_ctx)
|
||||
the text tokens
|
||||
xa : torch.Tensor, shape = (batch_size, n_mels, n_audio_ctx)
|
||||
the encoded audio features to be attended on
|
||||
include_embeddings : bool
|
||||
Whether to include intermediate values in the output to this function
|
||||
"""
|
||||
offset = next(iter(kv_cache.values())).shape[1] if kv_cache else 0
|
||||
x = self.token_embedding(x) + self.positional_embedding[offset : offset + x.shape[-1]]
|
||||
x = x.to(xa.dtype)
|
||||
|
||||
if include_embeddings:
|
||||
embeddings = [x.cpu().detach().numpy()]
|
||||
|
||||
for block in self.blocks:
|
||||
x = block(x, xa, mask=self.mask, kv_cache=kv_cache)
|
||||
if include_embeddings:
|
||||
embeddings.append(x.cpu().detach().numpy())
|
||||
|
||||
x = self.ln(x)
|
||||
logits = (x @ torch.transpose(self.token_embedding.weight.to(x.dtype), 0, 1)).float()
|
||||
|
||||
if include_embeddings:
|
||||
embeddings = np.stack(embeddings, axis=1)
|
||||
return logits, embeddings
|
||||
else:
|
||||
return logits
|
||||
|
||||
|
||||
class Whisper(nn.Module):
|
||||
def __init__(self, dims: ModelDimensions):
|
||||
super().__init__()
|
||||
self.dims = dims
|
||||
self.encoder = AudioEncoder(
|
||||
self.dims.n_mels,
|
||||
self.dims.n_audio_ctx,
|
||||
self.dims.n_audio_state,
|
||||
self.dims.n_audio_head,
|
||||
self.dims.n_audio_layer,
|
||||
)
|
||||
self.decoder = TextDecoder(
|
||||
self.dims.n_vocab,
|
||||
self.dims.n_text_ctx,
|
||||
self.dims.n_text_state,
|
||||
self.dims.n_text_head,
|
||||
self.dims.n_text_layer,
|
||||
)
|
||||
|
||||
def embed_audio(self, mel: torch.Tensor):
|
||||
return self.encoder.forward(mel)
|
||||
|
||||
def logits(self, tokens: torch.Tensor, audio_features: torch.Tensor):
|
||||
return self.decoder.forward(tokens, audio_features)
|
||||
|
||||
def forward(self, mel: torch.Tensor, tokens: torch.Tensor) -> Dict[str, torch.Tensor]:
|
||||
return self.decoder(tokens, self.encoder(mel))
|
||||
|
||||
@property
|
||||
def device(self):
|
||||
return next(self.parameters()).device
|
||||
|
||||
@property
|
||||
def is_multilingual(self):
|
||||
return self.dims.n_vocab == 51865
|
||||
|
||||
def install_kv_cache_hooks(self, cache: Optional[dict] = None):
|
||||
"""
|
||||
The `MultiHeadAttention` module optionally accepts `kv_cache` which stores the key and value
|
||||
tensors calculated for the previous positions. This method returns a dictionary that stores
|
||||
all caches, and the necessary hooks for the key and value projection modules that save the
|
||||
intermediate tensors to be reused during later calculations.
|
||||
|
||||
Returns
|
||||
-------
|
||||
cache : Dict[nn.Module, torch.Tensor]
|
||||
A dictionary object mapping the key/value projection modules to its cache
|
||||
hooks : List[RemovableHandle]
|
||||
List of PyTorch RemovableHandle objects to stop the hooks to be called
|
||||
"""
|
||||
cache = {**cache} if cache is not None else {}
|
||||
hooks = []
|
||||
|
||||
def save_to_cache(module, _, output):
|
||||
if module not in cache or output.shape[1] > self.decoder.positional_embedding.shape[0]:
|
||||
cache[module] = output # save as-is, for the first token or cross attention
|
||||
else:
|
||||
cache[module] = torch.cat([cache[module], output], dim=1).detach()
|
||||
return cache[module]
|
||||
|
||||
def install_hooks(layer: nn.Module):
|
||||
if isinstance(layer, MultiHeadAttention):
|
||||
hooks.append(layer.key.register_forward_hook(save_to_cache))
|
||||
hooks.append(layer.value.register_forward_hook(save_to_cache))
|
||||
|
||||
self.decoder.apply(install_hooks)
|
||||
return cache, hooks
|
||||
|
||||
detect_language = detect_language_function
|
||||
transcribe = transcribe_function
|
||||
decode = decode_function
|
||||
@@ -0,0 +1,2 @@
|
||||
from .basic import BasicTextNormalizer
|
||||
from .english import EnglishTextNormalizer
|
||||
@@ -0,0 +1,71 @@
|
||||
import re
|
||||
import unicodedata
|
||||
|
||||
import regex
|
||||
|
||||
# non-ASCII letters that are not separated by "NFKD" normalization
|
||||
ADDITIONAL_DIACRITICS = {
|
||||
"œ": "oe",
|
||||
"Œ": "OE",
|
||||
"ø": "o",
|
||||
"Ø": "O",
|
||||
"æ": "ae",
|
||||
"Æ": "AE",
|
||||
"ß": "ss",
|
||||
"ẞ": "SS",
|
||||
"đ": "d",
|
||||
"Đ": "D",
|
||||
"ð": "d",
|
||||
"Ð": "D",
|
||||
"þ": "th",
|
||||
"Þ": "th",
|
||||
"ł": "l",
|
||||
"Ł": "L",
|
||||
}
|
||||
|
||||
|
||||
def remove_symbols_and_diacritics(s: str, keep=""):
|
||||
"""
|
||||
Replace any other markers, symbols, and punctuations with a space,
|
||||
and drop any diacritics (category 'Mn' and some manual mappings)
|
||||
"""
|
||||
return "".join(
|
||||
c
|
||||
if c in keep
|
||||
else ADDITIONAL_DIACRITICS[c]
|
||||
if c in ADDITIONAL_DIACRITICS
|
||||
else ""
|
||||
if unicodedata.category(c) == "Mn"
|
||||
else " "
|
||||
if unicodedata.category(c)[0] in "MSP"
|
||||
else c
|
||||
for c in unicodedata.normalize("NFKD", s)
|
||||
)
|
||||
|
||||
|
||||
def remove_symbols(s: str):
|
||||
"""
|
||||
Replace any other markers, symbols, punctuations with a space, keeping diacritics
|
||||
"""
|
||||
return "".join(
|
||||
" " if unicodedata.category(c)[0] in "MSP" else c for c in unicodedata.normalize("NFKC", s)
|
||||
)
|
||||
|
||||
|
||||
class BasicTextNormalizer:
|
||||
def __init__(self, remove_diacritics: bool = False, split_letters: bool = False):
|
||||
self.clean = remove_symbols_and_diacritics if remove_diacritics else remove_symbols
|
||||
self.split_letters = split_letters
|
||||
|
||||
def __call__(self, s: str):
|
||||
s = s.lower()
|
||||
s = re.sub(r"[<\[][^>\]]*[>\]]", "", s) # remove words between brackets
|
||||
s = re.sub(r"\(([^)]+?)\)", "", s) # remove words between parenthesis
|
||||
s = self.clean(s).lower()
|
||||
|
||||
if self.split_letters:
|
||||
s = " ".join(regex.findall(r"\X", s, regex.U))
|
||||
|
||||
s = re.sub(r"\s+", " ", s) # replace any successive whitespace characters with a space
|
||||
|
||||
return s
|
||||
File diff suppressed because it is too large
Load Diff
@@ -0,0 +1,543 @@
|
||||
import json
|
||||
import os
|
||||
import re
|
||||
from fractions import Fraction
|
||||
from typing import Iterator, List, Match, Optional, Union
|
||||
|
||||
from more_itertools import windowed
|
||||
|
||||
from .basic import remove_symbols_and_diacritics
|
||||
|
||||
|
||||
class EnglishNumberNormalizer:
|
||||
"""
|
||||
Convert any spelled-out numbers into arabic numbers, while handling:
|
||||
|
||||
- remove any commas
|
||||
- keep the suffixes such as: `1960s`, `274th`, `32nd`, etc.
|
||||
- spell out currency symbols after the number. e.g. `$20 million` -> `20000000 dollars`
|
||||
- spell out `one` and `ones`
|
||||
- interpret successive single-digit numbers as nominal: `one oh one` -> `101`
|
||||
"""
|
||||
|
||||
def __init__(self):
|
||||
super().__init__()
|
||||
|
||||
self.zeros = {"o", "oh", "zero"}
|
||||
self.ones = {
|
||||
name: i
|
||||
for i, name in enumerate(
|
||||
[
|
||||
"one",
|
||||
"two",
|
||||
"three",
|
||||
"four",
|
||||
"five",
|
||||
"six",
|
||||
"seven",
|
||||
"eight",
|
||||
"nine",
|
||||
"ten",
|
||||
"eleven",
|
||||
"twelve",
|
||||
"thirteen",
|
||||
"fourteen",
|
||||
"fifteen",
|
||||
"sixteen",
|
||||
"seventeen",
|
||||
"eighteen",
|
||||
"nineteen",
|
||||
],
|
||||
start=1,
|
||||
)
|
||||
}
|
||||
self.ones_plural = {
|
||||
"sixes" if name == "six" else name + "s": (value, "s")
|
||||
for name, value in self.ones.items()
|
||||
}
|
||||
self.ones_ordinal = {
|
||||
"zeroth": (0, "th"),
|
||||
"first": (1, "st"),
|
||||
"second": (2, "nd"),
|
||||
"third": (3, "rd"),
|
||||
"fifth": (5, "th"),
|
||||
"twelfth": (12, "th"),
|
||||
**{
|
||||
name + ("h" if name.endswith("t") else "th"): (value, "th")
|
||||
for name, value in self.ones.items()
|
||||
if value > 3 and value != 5 and value != 12
|
||||
},
|
||||
}
|
||||
self.ones_suffixed = {**self.ones_plural, **self.ones_ordinal}
|
||||
|
||||
self.tens = {
|
||||
"twenty": 20,
|
||||
"thirty": 30,
|
||||
"forty": 40,
|
||||
"fifty": 50,
|
||||
"sixty": 60,
|
||||
"seventy": 70,
|
||||
"eighty": 80,
|
||||
"ninety": 90,
|
||||
}
|
||||
self.tens_plural = {
|
||||
name.replace("y", "ies"): (value, "s") for name, value in self.tens.items()
|
||||
}
|
||||
self.tens_ordinal = {
|
||||
name.replace("y", "ieth"): (value, "th") for name, value in self.tens.items()
|
||||
}
|
||||
self.tens_suffixed = {**self.tens_plural, **self.tens_ordinal}
|
||||
|
||||
self.multipliers = {
|
||||
"hundred": 100,
|
||||
"thousand": 1_000,
|
||||
"million": 1_000_000,
|
||||
"billion": 1_000_000_000,
|
||||
"trillion": 1_000_000_000_000,
|
||||
"quadrillion": 1_000_000_000_000_000,
|
||||
"quintillion": 1_000_000_000_000_000_000,
|
||||
"sextillion": 1_000_000_000_000_000_000_000,
|
||||
"septillion": 1_000_000_000_000_000_000_000_000,
|
||||
"octillion": 1_000_000_000_000_000_000_000_000_000,
|
||||
"nonillion": 1_000_000_000_000_000_000_000_000_000_000,
|
||||
"decillion": 1_000_000_000_000_000_000_000_000_000_000_000,
|
||||
}
|
||||
self.multipliers_plural = {
|
||||
name + "s": (value, "s") for name, value in self.multipliers.items()
|
||||
}
|
||||
self.multipliers_ordinal = {
|
||||
name + "th": (value, "th") for name, value in self.multipliers.items()
|
||||
}
|
||||
self.multipliers_suffixed = {**self.multipliers_plural, **self.multipliers_ordinal}
|
||||
self.decimals = {*self.ones, *self.tens, *self.zeros}
|
||||
|
||||
self.preceding_prefixers = {
|
||||
"minus": "-",
|
||||
"negative": "-",
|
||||
"plus": "+",
|
||||
"positive": "+",
|
||||
}
|
||||
self.following_prefixers = {
|
||||
"pound": "£",
|
||||
"pounds": "£",
|
||||
"euro": "€",
|
||||
"euros": "€",
|
||||
"dollar": "$",
|
||||
"dollars": "$",
|
||||
"cent": "¢",
|
||||
"cents": "¢",
|
||||
}
|
||||
self.prefixes = set(
|
||||
list(self.preceding_prefixers.values()) + list(self.following_prefixers.values())
|
||||
)
|
||||
self.suffixers = {
|
||||
"per": {"cent": "%"},
|
||||
"percent": "%",
|
||||
}
|
||||
self.specials = {"and", "double", "triple", "point"}
|
||||
|
||||
self.words = set(
|
||||
[
|
||||
key
|
||||
for mapping in [
|
||||
self.zeros,
|
||||
self.ones,
|
||||
self.ones_suffixed,
|
||||
self.tens,
|
||||
self.tens_suffixed,
|
||||
self.multipliers,
|
||||
self.multipliers_suffixed,
|
||||
self.preceding_prefixers,
|
||||
self.following_prefixers,
|
||||
self.suffixers,
|
||||
self.specials,
|
||||
]
|
||||
for key in mapping
|
||||
]
|
||||
)
|
||||
self.literal_words = {"one", "ones"}
|
||||
|
||||
def process_words(self, words: List[str]) -> Iterator[str]:
|
||||
prefix: Optional[str] = None
|
||||
value: Optional[Union[str, int]] = None
|
||||
skip = False
|
||||
|
||||
def to_fraction(s: str):
|
||||
try:
|
||||
return Fraction(s)
|
||||
except ValueError:
|
||||
return None
|
||||
|
||||
def output(result: Union[str, int]):
|
||||
nonlocal prefix, value
|
||||
result = str(result)
|
||||
if prefix is not None:
|
||||
result = prefix + result
|
||||
value = None
|
||||
prefix = None
|
||||
return result
|
||||
|
||||
if len(words) == 0:
|
||||
return
|
||||
|
||||
for prev, current, next in windowed([None] + words + [None], 3):
|
||||
if skip:
|
||||
skip = False
|
||||
continue
|
||||
|
||||
next_is_numeric = next is not None and re.match(r"^\d+(\.\d+)?$", next)
|
||||
has_prefix = current[0] in self.prefixes
|
||||
current_without_prefix = current[1:] if has_prefix else current
|
||||
if re.match(r"^\d+(\.\d+)?$", current_without_prefix):
|
||||
# arabic numbers (potentially with signs and fractions)
|
||||
f = to_fraction(current_without_prefix)
|
||||
assert f is not None
|
||||
if value is not None:
|
||||
if isinstance(value, str) and value.endswith("."):
|
||||
# concatenate decimals / ip address components
|
||||
value = str(value) + str(current)
|
||||
continue
|
||||
else:
|
||||
yield output(value)
|
||||
|
||||
prefix = current[0] if has_prefix else prefix
|
||||
if f.denominator == 1:
|
||||
value = f.numerator # store integers as int
|
||||
else:
|
||||
value = current_without_prefix
|
||||
elif current not in self.words:
|
||||
# non-numeric words
|
||||
if value is not None:
|
||||
yield output(value)
|
||||
yield output(current)
|
||||
elif current in self.zeros:
|
||||
value = str(value or "") + "0"
|
||||
elif current in self.ones:
|
||||
ones = self.ones[current]
|
||||
|
||||
if value is None:
|
||||
value = ones
|
||||
elif isinstance(value, str) or prev in self.ones:
|
||||
if prev in self.tens and ones < 10: # replace the last zero with the digit
|
||||
assert value[-1] == "0"
|
||||
value = value[:-1] + str(ones)
|
||||
else:
|
||||
value = str(value) + str(ones)
|
||||
elif ones < 10:
|
||||
if value % 10 == 0:
|
||||
value += ones
|
||||
else:
|
||||
value = str(value) + str(ones)
|
||||
else: # eleven to nineteen
|
||||
if value % 100 == 0:
|
||||
value += ones
|
||||
else:
|
||||
value = str(value) + str(ones)
|
||||
elif current in self.ones_suffixed:
|
||||
# ordinal or cardinal; yield the number right away
|
||||
ones, suffix = self.ones_suffixed[current]
|
||||
if value is None:
|
||||
yield output(str(ones) + suffix)
|
||||
elif isinstance(value, str) or prev in self.ones:
|
||||
if prev in self.tens and ones < 10:
|
||||
assert value[-1] == "0"
|
||||
yield output(value[:-1] + str(ones) + suffix)
|
||||
else:
|
||||
yield output(str(value) + str(ones) + suffix)
|
||||
elif ones < 10:
|
||||
if value % 10 == 0:
|
||||
yield output(str(value + ones) + suffix)
|
||||
else:
|
||||
yield output(str(value) + str(ones) + suffix)
|
||||
else: # eleven to nineteen
|
||||
if value % 100 == 0:
|
||||
yield output(str(value + ones) + suffix)
|
||||
else:
|
||||
yield output(str(value) + str(ones) + suffix)
|
||||
value = None
|
||||
elif current in self.tens:
|
||||
tens = self.tens[current]
|
||||
if value is None:
|
||||
value = tens
|
||||
elif isinstance(value, str):
|
||||
value = str(value) + str(tens)
|
||||
else:
|
||||
if value % 100 == 0:
|
||||
value += tens
|
||||
else:
|
||||
value = str(value) + str(tens)
|
||||
elif current in self.tens_suffixed:
|
||||
# ordinal or cardinal; yield the number right away
|
||||
tens, suffix = self.tens_suffixed[current]
|
||||
if value is None:
|
||||
yield output(str(tens) + suffix)
|
||||
elif isinstance(value, str):
|
||||
yield output(str(value) + str(tens) + suffix)
|
||||
else:
|
||||
if value % 100 == 0:
|
||||
yield output(str(value + tens) + suffix)
|
||||
else:
|
||||
yield output(str(value) + str(tens) + suffix)
|
||||
elif current in self.multipliers:
|
||||
multiplier = self.multipliers[current]
|
||||
if value is None:
|
||||
value = multiplier
|
||||
elif isinstance(value, str) or value == 0:
|
||||
f = to_fraction(value)
|
||||
p = f * multiplier if f is not None else None
|
||||
if f is not None and p.denominator == 1:
|
||||
value = p.numerator
|
||||
else:
|
||||
yield output(value)
|
||||
value = multiplier
|
||||
else:
|
||||
before = value // 1000 * 1000
|
||||
residual = value % 1000
|
||||
value = before + residual * multiplier
|
||||
elif current in self.multipliers_suffixed:
|
||||
multiplier, suffix = self.multipliers_suffixed[current]
|
||||
if value is None:
|
||||
yield output(str(multiplier) + suffix)
|
||||
elif isinstance(value, str):
|
||||
f = to_fraction(value)
|
||||
p = f * multiplier if f is not None else None
|
||||
if f is not None and p.denominator == 1:
|
||||
yield output(str(p.numerator) + suffix)
|
||||
else:
|
||||
yield output(value)
|
||||
yield output(str(multiplier) + suffix)
|
||||
else: # int
|
||||
before = value // 1000 * 1000
|
||||
residual = value % 1000
|
||||
value = before + residual * multiplier
|
||||
yield output(str(value) + suffix)
|
||||
value = None
|
||||
elif current in self.preceding_prefixers:
|
||||
# apply prefix (positive, minus, etc.) if it precedes a number
|
||||
if value is not None:
|
||||
yield output(value)
|
||||
|
||||
if next in self.words or next_is_numeric:
|
||||
prefix = self.preceding_prefixers[current]
|
||||
else:
|
||||
yield output(current)
|
||||
elif current in self.following_prefixers:
|
||||
# apply prefix (dollars, cents, etc.) only after a number
|
||||
if value is not None:
|
||||
prefix = self.following_prefixers[current]
|
||||
yield output(value)
|
||||
else:
|
||||
yield output(current)
|
||||
elif current in self.suffixers:
|
||||
# apply suffix symbols (percent -> '%')
|
||||
if value is not None:
|
||||
suffix = self.suffixers[current]
|
||||
if isinstance(suffix, dict):
|
||||
if next in suffix:
|
||||
yield output(str(value) + suffix[next])
|
||||
skip = True
|
||||
else:
|
||||
yield output(value)
|
||||
yield output(current)
|
||||
else:
|
||||
yield output(str(value) + suffix)
|
||||
else:
|
||||
yield output(current)
|
||||
elif current in self.specials:
|
||||
if next not in self.words and not next_is_numeric:
|
||||
# apply special handling only if the next word can be numeric
|
||||
if value is not None:
|
||||
yield output(value)
|
||||
yield output(current)
|
||||
elif current == "and":
|
||||
# ignore "and" after hundreds, thousands, etc.
|
||||
if prev not in self.multipliers:
|
||||
if value is not None:
|
||||
yield output(value)
|
||||
yield output(current)
|
||||
elif current == "double" or current == "triple":
|
||||
if next in self.ones or next in self.zeros:
|
||||
repeats = 2 if current == "double" else 3
|
||||
ones = self.ones.get(next, 0)
|
||||
value = str(value or "") + str(ones) * repeats
|
||||
skip = True
|
||||
else:
|
||||
if value is not None:
|
||||
yield output(value)
|
||||
yield output(current)
|
||||
elif current == "point":
|
||||
if next in self.decimals or next_is_numeric:
|
||||
value = str(value or "") + "."
|
||||
else:
|
||||
# should all have been covered at this point
|
||||
raise ValueError(f"Unexpected token: {current}")
|
||||
else:
|
||||
# all should have been covered at this point
|
||||
raise ValueError(f"Unexpected token: {current}")
|
||||
|
||||
if value is not None:
|
||||
yield output(value)
|
||||
|
||||
def preprocess(self, s: str):
|
||||
# replace "<number> and a half" with "<number> point five"
|
||||
results = []
|
||||
|
||||
segments = re.split(r"\band\s+a\s+half\b", s)
|
||||
for i, segment in enumerate(segments):
|
||||
if len(segment.strip()) == 0:
|
||||
continue
|
||||
if i == len(segments) - 1:
|
||||
results.append(segment)
|
||||
else:
|
||||
results.append(segment)
|
||||
last_word = segment.rsplit(maxsplit=2)[-1]
|
||||
if last_word in self.decimals or last_word in self.multipliers:
|
||||
results.append("point five")
|
||||
else:
|
||||
results.append("and a half")
|
||||
|
||||
s = " ".join(results)
|
||||
|
||||
# put a space at number/letter boundary
|
||||
s = re.sub(r"([a-z])([0-9])", r"\1 \2", s)
|
||||
s = re.sub(r"([0-9])([a-z])", r"\1 \2", s)
|
||||
|
||||
# but remove spaces which could be a suffix
|
||||
s = re.sub(r"([0-9])\s+(st|nd|rd|th|s)\b", r"\1\2", s)
|
||||
|
||||
return s
|
||||
|
||||
def postprocess(self, s: str):
|
||||
def combine_cents(m: Match):
|
||||
try:
|
||||
currency = m.group(1)
|
||||
integer = m.group(2)
|
||||
cents = int(m.group(3))
|
||||
return f"{currency}{integer}.{cents:02d}"
|
||||
except ValueError:
|
||||
return m.string
|
||||
|
||||
def extract_cents(m: Match):
|
||||
try:
|
||||
return f"¢{int(m.group(1))}"
|
||||
except ValueError:
|
||||
return m.string
|
||||
|
||||
# apply currency postprocessing; "$2 and ¢7" -> "$2.07"
|
||||
s = re.sub(r"([€£$])([0-9]+) (?:and )?¢([0-9]{1,2})\b", combine_cents, s)
|
||||
s = re.sub(r"[€£$]0.([0-9]{1,2})\b", extract_cents, s)
|
||||
|
||||
# write "one(s)" instead of "1(s)", just for the readability
|
||||
s = re.sub(r"\b1(s?)\b", r"one\1", s)
|
||||
|
||||
return s
|
||||
|
||||
def __call__(self, s: str):
|
||||
s = self.preprocess(s)
|
||||
s = " ".join(word for word in self.process_words(s.split()) if word is not None)
|
||||
s = self.postprocess(s)
|
||||
|
||||
return s
|
||||
|
||||
|
||||
class EnglishSpellingNormalizer:
|
||||
"""
|
||||
Applies British-American spelling mappings as listed in [1].
|
||||
|
||||
[1] https://www.tysto.com/uk-us-spelling-list.html
|
||||
"""
|
||||
|
||||
def __init__(self):
|
||||
mapping_path = os.path.join(os.path.dirname(__file__), "english.json")
|
||||
self.mapping = json.load(open(mapping_path))
|
||||
|
||||
def __call__(self, s: str):
|
||||
return " ".join(self.mapping.get(word, word) for word in s.split())
|
||||
|
||||
|
||||
class EnglishTextNormalizer:
|
||||
def __init__(self):
|
||||
self.ignore_patterns = r"\b(hmm|mm|mhm|mmm|uh|um)\b"
|
||||
self.replacers = {
|
||||
# common contractions
|
||||
r"\bwon't\b": "will not",
|
||||
r"\bcan't\b": "can not",
|
||||
r"\blet's\b": "let us",
|
||||
r"\bain't\b": "aint",
|
||||
r"\by'all\b": "you all",
|
||||
r"\bwanna\b": "want to",
|
||||
r"\bgotta\b": "got to",
|
||||
r"\bgonna\b": "going to",
|
||||
r"\bi'ma\b": "i am going to",
|
||||
r"\bimma\b": "i am going to",
|
||||
r"\bwoulda\b": "would have",
|
||||
r"\bcoulda\b": "could have",
|
||||
r"\bshoulda\b": "should have",
|
||||
r"\bma'am\b": "madam",
|
||||
# contractions in titles/prefixes
|
||||
r"\bmr\b": "mister ",
|
||||
r"\bmrs\b": "missus ",
|
||||
r"\bst\b": "saint ",
|
||||
r"\bdr\b": "doctor ",
|
||||
r"\bprof\b": "professor ",
|
||||
r"\bcapt\b": "captain ",
|
||||
r"\bgov\b": "governor ",
|
||||
r"\bald\b": "alderman ",
|
||||
r"\bgen\b": "general ",
|
||||
r"\bsen\b": "senator ",
|
||||
r"\brep\b": "representative ",
|
||||
r"\bpres\b": "president ",
|
||||
r"\brev\b": "reverend ",
|
||||
r"\bhon\b": "honorable ",
|
||||
r"\basst\b": "assistant ",
|
||||
r"\bassoc\b": "associate ",
|
||||
r"\blt\b": "lieutenant ",
|
||||
r"\bcol\b": "colonel ",
|
||||
r"\bjr\b": "junior ",
|
||||
r"\bsr\b": "senior ",
|
||||
r"\besq\b": "esquire ",
|
||||
# prefect tenses, ideally it should be any past participles, but it's harder..
|
||||
r"'d been\b": " had been",
|
||||
r"'s been\b": " has been",
|
||||
r"'d gone\b": " had gone",
|
||||
r"'s gone\b": " has gone",
|
||||
r"'d done\b": " had done", # "'s done" is ambiguous
|
||||
r"'s got\b": " has got",
|
||||
# general contractions
|
||||
r"n't\b": " not",
|
||||
r"'re\b": " are",
|
||||
r"'s\b": " is",
|
||||
r"'d\b": " would",
|
||||
r"'ll\b": " will",
|
||||
r"'t\b": " not",
|
||||
r"'ve\b": " have",
|
||||
r"'m\b": " am",
|
||||
}
|
||||
self.standardize_numbers = EnglishNumberNormalizer()
|
||||
self.standardize_spellings = EnglishSpellingNormalizer()
|
||||
|
||||
def __call__(self, s: str):
|
||||
s = s.lower()
|
||||
|
||||
s = re.sub(r"[<\[][^>\]]*[>\]]", "", s) # remove words between brackets
|
||||
s = re.sub(r"\(([^)]+?)\)", "", s) # remove words between parenthesis
|
||||
s = re.sub(self.ignore_patterns, "", s)
|
||||
s = re.sub(r"\s+'", "'", s) # standardize when there's a space before an apostrophe
|
||||
|
||||
for pattern, replacement in self.replacers.items():
|
||||
s = re.sub(pattern, replacement, s)
|
||||
|
||||
s = re.sub(r"(\d),(\d)", r"\1\2", s) # remove commas between digits
|
||||
s = re.sub(r"\.([^0-9]|$)", r" \1", s) # remove periods not followed by numbers
|
||||
s = remove_symbols_and_diacritics(s, keep=".%$¢€£") # keep some symbols for numerics
|
||||
|
||||
s = self.standardize_numbers(s)
|
||||
s = self.standardize_spellings(s)
|
||||
|
||||
# now remove prefix/suffix symbols that are not preceded/followed by numbers
|
||||
s = re.sub(r"[.$¢€£]([^0-9])", r" \1", s)
|
||||
s = re.sub(r"([^0-9])%", r"\1 ", s)
|
||||
|
||||
s = re.sub(r"\s+", " ", s) # replace any successive whitespace characters with a space
|
||||
|
||||
return s
|
||||
@@ -0,0 +1,331 @@
|
||||
import os
|
||||
from dataclasses import dataclass
|
||||
from functools import lru_cache
|
||||
from typing import List, Optional, Tuple, Union
|
||||
|
||||
import numpy as np
|
||||
import torch
|
||||
from transformers import GPT2TokenizerFast
|
||||
|
||||
LANGUAGES = {
|
||||
"en": "english",
|
||||
"zh": "chinese",
|
||||
"de": "german",
|
||||
"es": "spanish",
|
||||
"ru": "russian",
|
||||
"ko": "korean",
|
||||
"fr": "french",
|
||||
"ja": "japanese",
|
||||
"pt": "portuguese",
|
||||
"tr": "turkish",
|
||||
"pl": "polish",
|
||||
"ca": "catalan",
|
||||
"nl": "dutch",
|
||||
"ar": "arabic",
|
||||
"sv": "swedish",
|
||||
"it": "italian",
|
||||
"id": "indonesian",
|
||||
"hi": "hindi",
|
||||
"fi": "finnish",
|
||||
"vi": "vietnamese",
|
||||
"iw": "hebrew",
|
||||
"uk": "ukrainian",
|
||||
"el": "greek",
|
||||
"ms": "malay",
|
||||
"cs": "czech",
|
||||
"ro": "romanian",
|
||||
"da": "danish",
|
||||
"hu": "hungarian",
|
||||
"ta": "tamil",
|
||||
"no": "norwegian",
|
||||
"th": "thai",
|
||||
"ur": "urdu",
|
||||
"hr": "croatian",
|
||||
"bg": "bulgarian",
|
||||
"lt": "lithuanian",
|
||||
"la": "latin",
|
||||
"mi": "maori",
|
||||
"ml": "malayalam",
|
||||
"cy": "welsh",
|
||||
"sk": "slovak",
|
||||
"te": "telugu",
|
||||
"fa": "persian",
|
||||
"lv": "latvian",
|
||||
"bn": "bengali",
|
||||
"sr": "serbian",
|
||||
"az": "azerbaijani",
|
||||
"sl": "slovenian",
|
||||
"kn": "kannada",
|
||||
"et": "estonian",
|
||||
"mk": "macedonian",
|
||||
"br": "breton",
|
||||
"eu": "basque",
|
||||
"is": "icelandic",
|
||||
"hy": "armenian",
|
||||
"ne": "nepali",
|
||||
"mn": "mongolian",
|
||||
"bs": "bosnian",
|
||||
"kk": "kazakh",
|
||||
"sq": "albanian",
|
||||
"sw": "swahili",
|
||||
"gl": "galician",
|
||||
"mr": "marathi",
|
||||
"pa": "punjabi",
|
||||
"si": "sinhala",
|
||||
"km": "khmer",
|
||||
"sn": "shona",
|
||||
"yo": "yoruba",
|
||||
"so": "somali",
|
||||
"af": "afrikaans",
|
||||
"oc": "occitan",
|
||||
"ka": "georgian",
|
||||
"be": "belarusian",
|
||||
"tg": "tajik",
|
||||
"sd": "sindhi",
|
||||
"gu": "gujarati",
|
||||
"am": "amharic",
|
||||
"yi": "yiddish",
|
||||
"lo": "lao",
|
||||
"uz": "uzbek",
|
||||
"fo": "faroese",
|
||||
"ht": "haitian creole",
|
||||
"ps": "pashto",
|
||||
"tk": "turkmen",
|
||||
"nn": "nynorsk",
|
||||
"mt": "maltese",
|
||||
"sa": "sanskrit",
|
||||
"lb": "luxembourgish",
|
||||
"my": "myanmar",
|
||||
"bo": "tibetan",
|
||||
"tl": "tagalog",
|
||||
"mg": "malagasy",
|
||||
"as": "assamese",
|
||||
"tt": "tatar",
|
||||
"haw": "hawaiian",
|
||||
"ln": "lingala",
|
||||
"ha": "hausa",
|
||||
"ba": "bashkir",
|
||||
"jw": "javanese",
|
||||
"su": "sundanese",
|
||||
}
|
||||
|
||||
# language code lookup by name, with a few language aliases
|
||||
TO_LANGUAGE_CODE = {
|
||||
**{language: code for code, language in LANGUAGES.items()},
|
||||
"burmese": "my",
|
||||
"valencian": "ca",
|
||||
"flemish": "nl",
|
||||
"haitian": "ht",
|
||||
"letzeburgesch": "lb",
|
||||
"pushto": "ps",
|
||||
"panjabi": "pa",
|
||||
"moldavian": "ro",
|
||||
"moldovan": "ro",
|
||||
"sinhalese": "si",
|
||||
"castilian": "es",
|
||||
}
|
||||
|
||||
|
||||
@dataclass(frozen=True)
|
||||
class Tokenizer:
|
||||
"""A thin wrapper around `GPT2TokenizerFast` providing quick access to special tokens"""
|
||||
|
||||
tokenizer: "GPT2TokenizerFast"
|
||||
language: Optional[str]
|
||||
sot_sequence: Tuple[int]
|
||||
|
||||
def encode(self, text, **kwargs):
|
||||
return self.tokenizer.encode(text, **kwargs)
|
||||
|
||||
def decode(self, token_ids: Union[int, List[int], np.ndarray, torch.Tensor], **kwargs):
|
||||
return self.tokenizer.decode(token_ids, **kwargs)
|
||||
|
||||
def decode_with_timestamps(self, tokens) -> str:
|
||||
"""
|
||||
Timestamp tokens are above the special tokens' id range and are ignored by `decode()`.
|
||||
This method decodes given tokens with timestamps tokens annotated, e.g. "<|1.08|>".
|
||||
"""
|
||||
outputs = [[]]
|
||||
for token in tokens:
|
||||
if token >= self.timestamp_begin:
|
||||
timestamp = f"<|{(token - self.timestamp_begin) * 0.02:.2f}|>"
|
||||
outputs.append(timestamp)
|
||||
outputs.append([])
|
||||
else:
|
||||
outputs[-1].append(token)
|
||||
outputs = [s if isinstance(s, str) else self.tokenizer.decode(s) for s in outputs]
|
||||
return "".join(outputs)
|
||||
|
||||
@property
|
||||
@lru_cache()
|
||||
def eot(self) -> int:
|
||||
return self.tokenizer.eos_token_id
|
||||
|
||||
@property
|
||||
@lru_cache()
|
||||
def sot(self) -> int:
|
||||
return self._get_single_token_id("<|startoftranscript|>")
|
||||
|
||||
@property
|
||||
@lru_cache()
|
||||
def sot_lm(self) -> int:
|
||||
return self._get_single_token_id("<|startoflm|>")
|
||||
|
||||
@property
|
||||
@lru_cache()
|
||||
def sot_prev(self) -> int:
|
||||
return self._get_single_token_id("<|startofprev|>")
|
||||
|
||||
@property
|
||||
@lru_cache()
|
||||
def no_speech(self) -> int:
|
||||
return self._get_single_token_id("<|nospeech|>")
|
||||
|
||||
@property
|
||||
@lru_cache()
|
||||
def no_timestamps(self) -> int:
|
||||
return self._get_single_token_id("<|notimestamps|>")
|
||||
|
||||
@property
|
||||
@lru_cache()
|
||||
def timestamp_begin(self) -> int:
|
||||
return self.tokenizer.all_special_ids[-1] + 1
|
||||
|
||||
@property
|
||||
@lru_cache()
|
||||
def language_token(self) -> int:
|
||||
"""Returns the token id corresponding to the value of the `language` field"""
|
||||
if self.language is None:
|
||||
raise ValueError(f"This tokenizer does not have language token configured")
|
||||
|
||||
additional_tokens = dict(
|
||||
zip(
|
||||
self.tokenizer.additional_special_tokens,
|
||||
self.tokenizer.additional_special_tokens_ids,
|
||||
)
|
||||
)
|
||||
candidate = f"<|{self.language}|>"
|
||||
if candidate in additional_tokens:
|
||||
return additional_tokens[candidate]
|
||||
|
||||
raise KeyError(f"Language {self.language} not found in tokenizer.")
|
||||
|
||||
@property
|
||||
@lru_cache()
|
||||
def all_language_tokens(self) -> Tuple[int]:
|
||||
result = []
|
||||
for token, token_id in zip(
|
||||
self.tokenizer.additional_special_tokens,
|
||||
self.tokenizer.additional_special_tokens_ids,
|
||||
):
|
||||
if token.strip("<|>") in LANGUAGES:
|
||||
result.append(token_id)
|
||||
return tuple(result)
|
||||
|
||||
@property
|
||||
@lru_cache()
|
||||
def all_language_codes(self) -> Tuple[str]:
|
||||
return tuple(self.decode([l]).strip("<|>") for l in self.all_language_tokens)
|
||||
|
||||
@property
|
||||
@lru_cache()
|
||||
def sot_sequence_including_notimestamps(self) -> Tuple[int]:
|
||||
return tuple(list(self.sot_sequence) + [self.no_timestamps])
|
||||
|
||||
@property
|
||||
@lru_cache()
|
||||
def non_speech_tokens(self) -> Tuple[int]:
|
||||
"""
|
||||
Returns the list of tokens to suppress in order to avoid any speaker tags or non-speech
|
||||
annotations, to prevent sampling texts that are not actually spoken in the audio, e.g.
|
||||
|
||||
- ♪♪♪
|
||||
- ( SPEAKING FOREIGN LANGUAGE )
|
||||
- [DAVID] Hey there,
|
||||
|
||||
keeping basic punctuations like commas, periods, question marks, exclamation points, etc.
|
||||
"""
|
||||
symbols = list("\"#()*+/:;<=>@[\\]^_`{|}~「」『』")
|
||||
symbols += "<< >> <<< >>> -- --- -( -[ (' (\" (( )) ((( ))) [[ ]] {{ }} ♪♪ ♪♪♪".split()
|
||||
|
||||
# symbols that may be a single token or multiple tokens depending on the tokenizer.
|
||||
# In case they're multiple tokens, suppress the first token, which is safe because:
|
||||
# These are between U+2640 and U+267F miscellaneous symbols that are okay to suppress
|
||||
# in generations, and in the 3-byte UTF-8 representation they share the first two bytes.
|
||||
miscellaneous = set("♩♪♫♬♭♮♯")
|
||||
assert all(0x2640 <= ord(c) <= 0x267F for c in miscellaneous)
|
||||
|
||||
# allow hyphens "-" and single quotes "'" between words, but not at the beginning of a word
|
||||
result = {self.tokenizer.encode(" -")[0], self.tokenizer.encode(" '")[0]}
|
||||
for symbol in symbols + list(miscellaneous):
|
||||
for tokens in [self.tokenizer.encode(symbol), self.tokenizer.encode(" " + symbol)]:
|
||||
if len(tokens) == 1 or symbol in miscellaneous:
|
||||
result.add(tokens[0])
|
||||
|
||||
return tuple(sorted(result))
|
||||
|
||||
def _get_single_token_id(self, text) -> int:
|
||||
tokens = self.tokenizer.encode(text)
|
||||
assert len(tokens) == 1, f"{text} is not encoded as a single token"
|
||||
return tokens[0]
|
||||
|
||||
|
||||
@lru_cache(maxsize=None)
|
||||
def build_tokenizer(name: str = "gpt2"):
|
||||
os.environ["TOKENIZERS_PARALLELISM"] = "false"
|
||||
path = os.path.join(os.path.dirname(__file__), "assets", name)
|
||||
tokenizer = GPT2TokenizerFast.from_pretrained(path)
|
||||
|
||||
specials = [
|
||||
"<|startoftranscript|>",
|
||||
*[f"<|{lang}|>" for lang in LANGUAGES.keys()],
|
||||
"<|translate|>",
|
||||
"<|transcribe|>",
|
||||
"<|startoflm|>",
|
||||
"<|startofprev|>",
|
||||
"<|nospeech|>",
|
||||
"<|notimestamps|>",
|
||||
]
|
||||
|
||||
tokenizer.add_special_tokens(dict(additional_special_tokens=specials))
|
||||
return tokenizer
|
||||
|
||||
|
||||
@lru_cache(maxsize=None)
|
||||
def get_tokenizer(
|
||||
multilingual: bool,
|
||||
*,
|
||||
task: Optional[str] = None, # Literal["transcribe", "translate", None]
|
||||
language: Optional[str] = None,
|
||||
) -> Tokenizer:
|
||||
if language is not None:
|
||||
language = language.lower()
|
||||
if language not in LANGUAGES:
|
||||
if language in TO_LANGUAGE_CODE:
|
||||
language = TO_LANGUAGE_CODE[language]
|
||||
else:
|
||||
raise ValueError(f"Unsupported language: {language}")
|
||||
|
||||
if multilingual:
|
||||
tokenizer_name = "multilingual"
|
||||
task = task or "transcribe"
|
||||
language = language or "en"
|
||||
else:
|
||||
tokenizer_name = "gpt2"
|
||||
task = None
|
||||
language = None
|
||||
|
||||
tokenizer = build_tokenizer(name=tokenizer_name)
|
||||
all_special_ids: List[int] = tokenizer.all_special_ids
|
||||
sot: int = all_special_ids[1]
|
||||
translate: int = all_special_ids[-6]
|
||||
transcribe: int = all_special_ids[-5]
|
||||
|
||||
langs = tuple(LANGUAGES.keys())
|
||||
sot_sequence = [sot]
|
||||
if language is not None:
|
||||
sot_sequence.append(sot + 1 + langs.index(language))
|
||||
if task is not None:
|
||||
sot_sequence.append(transcribe if task == "transcribe" else translate)
|
||||
|
||||
return Tokenizer(tokenizer=tokenizer, language=language, sot_sequence=tuple(sot_sequence))
|
||||
@@ -0,0 +1,207 @@
|
||||
import argparse
|
||||
import os
|
||||
import warnings
|
||||
from typing import List, Optional, Tuple, Union, TYPE_CHECKING
|
||||
|
||||
import numpy as np
|
||||
import torch
|
||||
import tqdm
|
||||
|
||||
from .audio import SAMPLE_RATE, N_FRAMES, HOP_LENGTH, pad_or_trim, log_mel_spectrogram
|
||||
from .decoding import DecodingOptions, DecodingResult
|
||||
from .tokenizer import LANGUAGES, TO_LANGUAGE_CODE, get_tokenizer
|
||||
from .utils import exact_div, format_timestamp, optional_int, optional_float, str2bool, write_txt, write_vtt, write_srt
|
||||
|
||||
if TYPE_CHECKING:
|
||||
from .model import Whisper
|
||||
|
||||
|
||||
def transcribe(
|
||||
model: "Whisper",
|
||||
audio: Union[str, np.ndarray, torch.Tensor],
|
||||
*,
|
||||
verbose: Optional[bool] = None,
|
||||
temperature: Union[float, Tuple[float, ...]] = (0.0, 0.2, 0.4, 0.6, 0.8, 1.0),
|
||||
compression_ratio_threshold: Optional[float] = 2.4,
|
||||
logprob_threshold: Optional[float] = -1.0,
|
||||
no_speech_threshold: Optional[float] = 0.6,
|
||||
condition_on_previous_text: bool = True,
|
||||
force_extraction: bool = False,
|
||||
**decode_options,
|
||||
):
|
||||
"""
|
||||
Transcribe an audio file using Whisper
|
||||
|
||||
Parameters
|
||||
----------
|
||||
model: Whisper
|
||||
The Whisper model instance
|
||||
|
||||
audio: Union[str, np.ndarray, torch.Tensor]
|
||||
The path to the audio file to open, or the audio waveform
|
||||
|
||||
verbose: bool
|
||||
Whether to display the text being decoded to the console. If True, displays all the details,
|
||||
If False, displays minimal details. If None, does not display anything
|
||||
|
||||
temperature: Union[float, Tuple[float, ...]]
|
||||
Temperature for sampling. It can be a tuple of temperatures, which will be successfully used
|
||||
upon failures according to either `compression_ratio_threshold` or `logprob_threshold`.
|
||||
|
||||
compression_ratio_threshold: float
|
||||
If the gzip compression ratio is above this value, treat as failed
|
||||
|
||||
logprob_threshold: float
|
||||
If the average log probability over sampled tokens is below this value, treat as failed
|
||||
|
||||
no_speech_threshold: float
|
||||
If the no_speech probability is higher than this value AND the average log probability
|
||||
over sampled tokens is below `logprob_threshold`, consider the segment as silent
|
||||
|
||||
condition_on_previous_text: bool
|
||||
if True, the previous output of the model is provided as a prompt for the next window;
|
||||
disabling may make the text inconsistent across windows, but the model becomes less prone to
|
||||
getting stuck in a failure loop, such as repetition looping or timestamps going out of sync.
|
||||
|
||||
decode_options: dict
|
||||
Keyword arguments to construct `DecodingOptions` instances
|
||||
|
||||
Returns
|
||||
-------
|
||||
A dictionary containing the resulting text ("text") and segment-level details ("segments"), and
|
||||
the spoken language ("language"), which is detected when `decode_options["language"]` is None.
|
||||
"""
|
||||
dtype = torch.float16 if decode_options.get("fp16", True) else torch.float32
|
||||
if model.device == torch.device("cpu"):
|
||||
if torch.cuda.is_available():
|
||||
warnings.warn("Performing inference on CPU when CUDA is available")
|
||||
if dtype == torch.float16:
|
||||
warnings.warn("FP16 is not supported on CPU; using FP32 instead")
|
||||
dtype = torch.float32
|
||||
|
||||
if dtype == torch.float32:
|
||||
decode_options["fp16"] = False
|
||||
|
||||
mel = log_mel_spectrogram(audio)
|
||||
|
||||
all_segments = []
|
||||
def add_segment(
|
||||
*, start: float, end: float, encoder_embeddings
|
||||
):
|
||||
|
||||
all_segments.append(
|
||||
{
|
||||
"start": start,
|
||||
"end": end,
|
||||
"encoder_embeddings":encoder_embeddings,
|
||||
}
|
||||
)
|
||||
# show the progress bar when verbose is False (otherwise the transcribed text will be printed)
|
||||
num_frames = mel.shape[-1]
|
||||
seek = 0
|
||||
previous_seek_value = seek
|
||||
sample_skip = 3000 #
|
||||
with tqdm.tqdm(total=num_frames, unit='frames', disable=verbose is not False) as pbar:
|
||||
while seek < num_frames:
|
||||
# seek是开始的帧数
|
||||
end_seek = min(seek + sample_skip, num_frames)
|
||||
segment = pad_or_trim(mel[:,seek:seek+sample_skip], N_FRAMES).to(model.device).to(dtype)
|
||||
|
||||
single = segment.ndim == 2
|
||||
if single:
|
||||
segment = segment.unsqueeze(0)
|
||||
if dtype == torch.float16:
|
||||
segment = segment.half()
|
||||
audio_features, embeddings = model.encoder(segment, include_embeddings = True)
|
||||
|
||||
encoder_embeddings = embeddings
|
||||
#print(f"encoder_embeddings shape {encoder_embeddings.shape}")
|
||||
add_segment(
|
||||
start=seek,
|
||||
end=end_seek,
|
||||
#text_tokens=tokens,
|
||||
#result=result,
|
||||
encoder_embeddings=encoder_embeddings,
|
||||
)
|
||||
seek+=sample_skip
|
||||
|
||||
return dict(segments=all_segments)
|
||||
|
||||
|
||||
def cli():
|
||||
from . import available_models
|
||||
|
||||
parser = argparse.ArgumentParser(formatter_class=argparse.ArgumentDefaultsHelpFormatter)
|
||||
parser.add_argument("audio", nargs="+", type=str, help="audio file(s) to transcribe")
|
||||
parser.add_argument("--model", default="small", choices=available_models(), help="name of the Whisper model to use")
|
||||
parser.add_argument("--model_dir", type=str, default=None, help="the path to save model files; uses ~/.cache/whisper by default")
|
||||
parser.add_argument("--device", default="cuda" if torch.cuda.is_available() else "cpu", help="device to use for PyTorch inference")
|
||||
parser.add_argument("--output_dir", "-o", type=str, default=".", help="directory to save the outputs")
|
||||
parser.add_argument("--verbose", type=str2bool, default=True, help="whether to print out the progress and debug messages")
|
||||
|
||||
parser.add_argument("--task", type=str, default="transcribe", choices=["transcribe", "translate"], help="whether to perform X->X speech recognition ('transcribe') or X->English translation ('translate')")
|
||||
parser.add_argument("--language", type=str, default=None, choices=sorted(LANGUAGES.keys()) + sorted([k.title() for k in TO_LANGUAGE_CODE.keys()]), help="language spoken in the audio, specify None to perform language detection")
|
||||
|
||||
parser.add_argument("--temperature", type=float, default=0, help="temperature to use for sampling")
|
||||
parser.add_argument("--best_of", type=optional_int, default=5, help="number of candidates when sampling with non-zero temperature")
|
||||
parser.add_argument("--beam_size", type=optional_int, default=5, help="number of beams in beam search, only applicable when temperature is zero")
|
||||
parser.add_argument("--patience", type=float, default=None, help="optional patience value to use in beam decoding, as in https://arxiv.org/abs/2204.05424, the default (1.0) is equivalent to conventional beam search")
|
||||
parser.add_argument("--length_penalty", type=float, default=None, help="optional token length penalty coefficient (alpha) as in https://arxiv.org/abs/1609.08144, uses simple length normalization by default")
|
||||
|
||||
parser.add_argument("--suppress_tokens", type=str, default="-1", help="comma-separated list of token ids to suppress during sampling; '-1' will suppress most special characters except common punctuations")
|
||||
parser.add_argument("--initial_prompt", type=str, default=None, help="optional text to provide as a prompt for the first window.")
|
||||
parser.add_argument("--condition_on_previous_text", type=str2bool, default=True, help="if True, provide the previous output of the model as a prompt for the next window; disabling may make the text inconsistent across windows, but the model becomes less prone to getting stuck in a failure loop")
|
||||
parser.add_argument("--fp16", type=str2bool, default=True, help="whether to perform inference in fp16; True by default")
|
||||
|
||||
parser.add_argument("--temperature_increment_on_fallback", type=optional_float, default=0.2, help="temperature to increase when falling back when the decoding fails to meet either of the thresholds below")
|
||||
parser.add_argument("--compression_ratio_threshold", type=optional_float, default=2.4, help="if the gzip compression ratio is higher than this value, treat the decoding as failed")
|
||||
parser.add_argument("--logprob_threshold", type=optional_float, default=-1.0, help="if the average log probability is lower than this value, treat the decoding as failed")
|
||||
parser.add_argument("--no_speech_threshold", type=optional_float, default=0.6, help="if the probability of the <|nospeech|> token is higher than this value AND the decoding has failed due to `logprob_threshold`, consider the segment as silence")
|
||||
parser.add_argument("--threads", type=optional_int, default=0, help="number of threads used by torch for CPU inference; supercedes MKL_NUM_THREADS/OMP_NUM_THREADS")
|
||||
|
||||
args = parser.parse_args().__dict__
|
||||
model_name: str = args.pop("model")
|
||||
model_dir: str = args.pop("model_dir")
|
||||
output_dir: str = args.pop("output_dir")
|
||||
device: str = args.pop("device")
|
||||
os.makedirs(output_dir, exist_ok=True)
|
||||
|
||||
if model_name.endswith(".en") and args["language"] not in {"en", "English"}:
|
||||
if args["language"] is not None:
|
||||
warnings.warn(f"{model_name} is an English-only model but receipted '{args['language']}'; using English instead.")
|
||||
args["language"] = "en"
|
||||
|
||||
temperature = args.pop("temperature")
|
||||
temperature_increment_on_fallback = args.pop("temperature_increment_on_fallback")
|
||||
if temperature_increment_on_fallback is not None:
|
||||
temperature = tuple(np.arange(temperature, 1.0 + 1e-6, temperature_increment_on_fallback))
|
||||
else:
|
||||
temperature = [temperature]
|
||||
|
||||
threads = args.pop("threads")
|
||||
if threads > 0:
|
||||
torch.set_num_threads(threads)
|
||||
|
||||
from . import load_model
|
||||
model = load_model(model_name, device=device, download_root=model_dir)
|
||||
|
||||
for audio_path in args.pop("audio"):
|
||||
result = transcribe(model, audio_path, temperature=temperature, **args)
|
||||
|
||||
audio_basename = os.path.basename(audio_path)
|
||||
|
||||
# save TXT
|
||||
with open(os.path.join(output_dir, audio_basename + ".txt"), "w", encoding="utf-8") as txt:
|
||||
write_txt(result["segments"], file=txt)
|
||||
|
||||
# save VTT
|
||||
with open(os.path.join(output_dir, audio_basename + ".vtt"), "w", encoding="utf-8") as vtt:
|
||||
write_vtt(result["segments"], file=vtt)
|
||||
|
||||
# save SRT
|
||||
with open(os.path.join(output_dir, audio_basename + ".srt"), "w", encoding="utf-8") as srt:
|
||||
write_srt(result["segments"], file=srt)
|
||||
|
||||
|
||||
if __name__ == '__main__':
|
||||
cli()
|
||||
@@ -0,0 +1,87 @@
|
||||
import zlib
|
||||
from typing import Iterator, TextIO
|
||||
|
||||
|
||||
def exact_div(x, y):
|
||||
assert x % y == 0
|
||||
return x // y
|
||||
|
||||
|
||||
def str2bool(string):
|
||||
str2val = {"True": True, "False": False}
|
||||
if string in str2val:
|
||||
return str2val[string]
|
||||
else:
|
||||
raise ValueError(f"Expected one of {set(str2val.keys())}, got {string}")
|
||||
|
||||
|
||||
def optional_int(string):
|
||||
return None if string == "None" else int(string)
|
||||
|
||||
|
||||
def optional_float(string):
|
||||
return None if string == "None" else float(string)
|
||||
|
||||
|
||||
def compression_ratio(text) -> float:
|
||||
return len(text) / len(zlib.compress(text.encode("utf-8")))
|
||||
|
||||
|
||||
def format_timestamp(seconds: float, always_include_hours: bool = False, decimal_marker: str = '.'):
|
||||
assert seconds >= 0, "non-negative timestamp expected"
|
||||
milliseconds = round(seconds * 1000.0)
|
||||
|
||||
hours = milliseconds // 3_600_000
|
||||
milliseconds -= hours * 3_600_000
|
||||
|
||||
minutes = milliseconds // 60_000
|
||||
milliseconds -= minutes * 60_000
|
||||
|
||||
seconds = milliseconds // 1_000
|
||||
milliseconds -= seconds * 1_000
|
||||
|
||||
hours_marker = f"{hours:02d}:" if always_include_hours or hours > 0 else ""
|
||||
return f"{hours_marker}{minutes:02d}:{seconds:02d}{decimal_marker}{milliseconds:03d}"
|
||||
|
||||
|
||||
def write_txt(transcript: Iterator[dict], file: TextIO):
|
||||
for segment in transcript:
|
||||
print(segment['text'].strip(), file=file, flush=True)
|
||||
|
||||
|
||||
def write_vtt(transcript: Iterator[dict], file: TextIO):
|
||||
print("WEBVTT\n", file=file)
|
||||
for segment in transcript:
|
||||
print(
|
||||
f"{format_timestamp(segment['start'])} --> {format_timestamp(segment['end'])}\n"
|
||||
f"{segment['text'].strip().replace('-->', '->')}\n",
|
||||
file=file,
|
||||
flush=True,
|
||||
)
|
||||
|
||||
|
||||
def write_srt(transcript: Iterator[dict], file: TextIO):
|
||||
"""
|
||||
Write a transcript to a file in SRT format.
|
||||
|
||||
Example usage:
|
||||
from pathlib import Path
|
||||
from whisper.utils import write_srt
|
||||
|
||||
result = transcribe(model, audio_path, temperature=temperature, **args)
|
||||
|
||||
# save SRT
|
||||
audio_basename = Path(audio_path).stem
|
||||
with open(Path(output_dir) / (audio_basename + ".srt"), "w", encoding="utf-8") as srt:
|
||||
write_srt(result["segments"], file=srt)
|
||||
"""
|
||||
for i, segment in enumerate(transcript, start=1):
|
||||
# write srt lines
|
||||
print(
|
||||
f"{i}\n"
|
||||
f"{format_timestamp(segment['start'], always_include_hours=True, decimal_marker=',')} --> "
|
||||
f"{format_timestamp(segment['end'], always_include_hours=True, decimal_marker=',')}\n"
|
||||
f"{segment['text'].strip().replace('-->', '->')}\n",
|
||||
file=file,
|
||||
flush=True,
|
||||
)
|
||||
@@ -0,0 +1,101 @@
|
||||
# TODO: Adapted from cli
|
||||
from typing import Callable, List, Optional
|
||||
|
||||
import numpy as np
|
||||
|
||||
|
||||
def ordered_halving(val):
|
||||
bin_str = f"{val:064b}"
|
||||
bin_flip = bin_str[::-1]
|
||||
as_int = int(bin_flip, 2)
|
||||
|
||||
return as_int / (1 << 64)
|
||||
|
||||
|
||||
def uniform(
|
||||
step: int = ...,
|
||||
num_steps: Optional[int] = None,
|
||||
num_frames: int = ...,
|
||||
context_size: Optional[int] = None,
|
||||
context_stride: int = 3,
|
||||
context_overlap: int = 4,
|
||||
closed_loop: bool = True,
|
||||
):
|
||||
if num_frames <= context_size:
|
||||
yield list(range(num_frames))
|
||||
return
|
||||
|
||||
context_stride = min(
|
||||
context_stride, int(np.ceil(np.log2(num_frames / context_size))) + 1
|
||||
)
|
||||
|
||||
for context_step in 1 << np.arange(context_stride):
|
||||
pad = int(round(num_frames * ordered_halving(step)))
|
||||
for j in range(
|
||||
int(ordered_halving(step) * context_step) + pad,
|
||||
num_frames + pad + (0 if closed_loop else -context_overlap),
|
||||
(context_size * context_step - context_overlap),
|
||||
):
|
||||
yield [
|
||||
e % num_frames
|
||||
for e in range(j, j + context_size * context_step, context_step)
|
||||
]
|
||||
|
||||
|
||||
def get_context_scheduler(name: str) -> Callable:
|
||||
if name == "uniform":
|
||||
return uniform
|
||||
else:
|
||||
raise ValueError(f"Unknown context_overlap policy {name}")
|
||||
|
||||
|
||||
def get_total_steps(
|
||||
scheduler,
|
||||
timesteps: List[int],
|
||||
num_steps: Optional[int] = None,
|
||||
num_frames: int = ...,
|
||||
context_size: Optional[int] = None,
|
||||
context_stride: int = 3,
|
||||
context_overlap: int = 4,
|
||||
closed_loop: bool = True,
|
||||
):
|
||||
return sum(
|
||||
len(
|
||||
list(
|
||||
scheduler(
|
||||
i,
|
||||
num_steps,
|
||||
num_frames,
|
||||
context_size,
|
||||
context_stride,
|
||||
context_overlap,
|
||||
)
|
||||
)
|
||||
)
|
||||
for i in range(len(timesteps))
|
||||
)
|
||||
|
||||
def modify_context_all(context):
|
||||
'''
|
||||
modify this context to remove first indexes in last context and expand it to 16 frames
|
||||
example :
|
||||
input : [[240, 241, 242, 243, 244, 245, 246, 247, 248, 249, 0, 1, 2, 3, 4, 5]]
|
||||
output : s = [[234, 235, 236, 237, 238, 239, 240, 241, 242, 243, 244, 245, 246, 247, 248, 249]]
|
||||
'''
|
||||
context = context[0]
|
||||
|
||||
if context[0]==0:
|
||||
len_context = len(context)
|
||||
context_new = [list(np.arange(0, 12))]
|
||||
|
||||
else:
|
||||
len_context = 1
|
||||
max_value_list = context[0]
|
||||
for c, value in enumerate(context):
|
||||
if value > max_value_list:
|
||||
len_context = c + 1
|
||||
max_value_list = value
|
||||
|
||||
context_new = [list(np.arange(max(0,max_value_list-12+1), max_value_list+1))]
|
||||
|
||||
return context_new, len_context
|
||||
@@ -0,0 +1,487 @@
|
||||
# pylint: disable=R0801
|
||||
"""
|
||||
This module is responsible for animating faces in videos using a combination of deep learning techniques.
|
||||
It provides a pipeline for generating face animations by processing video frames and extracting face features.
|
||||
The module utilizes various schedulers and utilities for efficient face animation and supports different types
|
||||
of latents for more control over the animation process.
|
||||
|
||||
Functions and Classes:
|
||||
- FaceAnimatePipeline: A class that extends the DiffusionPipeline class from the diffusers library to handle face animation tasks.
|
||||
- __init__: Initializes the pipeline with the necessary components (VAE, UNets, face locator, etc.).
|
||||
- prepare_latents: Generates or loads latents for the animation process, scaling them according to the scheduler's requirements.
|
||||
- prepare_extra_step_kwargs: Prepares extra keyword arguments for the scheduler step, ensuring compatibility with different schedulers.
|
||||
- decode_latents: Decodes the latents into video frames, ready for animation.
|
||||
|
||||
Usage:
|
||||
- Import the necessary packages and classes.
|
||||
- Create a FaceAnimatePipeline instance with the required components.
|
||||
- Prepare the latents for the animation process.
|
||||
- Use the pipeline to generate the animated video.
|
||||
|
||||
Note:
|
||||
- This module is designed to work with the diffusers library, which provides the underlying framework for face animation using deep learning.
|
||||
- The module is intended for research and development purposes, and further optimization and customization may be required for specific use cases.
|
||||
"""
|
||||
|
||||
import inspect
|
||||
from dataclasses import dataclass
|
||||
from typing import Callable, List, Optional, Union
|
||||
import math
|
||||
import numpy as np
|
||||
import torch
|
||||
from diffusers import (DDIMScheduler, DiffusionPipeline,
|
||||
DPMSolverMultistepScheduler,
|
||||
EulerAncestralDiscreteScheduler, EulerDiscreteScheduler,
|
||||
LMSDiscreteScheduler, PNDMScheduler)
|
||||
from diffusers.image_processor import VaeImageProcessor
|
||||
from diffusers.utils import BaseOutput
|
||||
from diffusers.utils.torch_utils import randn_tensor
|
||||
from einops import rearrange, repeat
|
||||
from tqdm import tqdm
|
||||
|
||||
from .context import get_context_scheduler, modify_context_all
|
||||
from ..models.mutual_self_attention import ReferenceAttentionControl
|
||||
|
||||
|
||||
@dataclass
|
||||
class FaceAnimatePipelineOutput(BaseOutput):
|
||||
"""
|
||||
FaceAnimatePipelineOutput is a custom class that inherits from BaseOutput and represents the output of the FaceAnimatePipeline.
|
||||
|
||||
Attributes:
|
||||
videos (Union[torch.Tensor, np.ndarray]): A tensor or numpy array containing the generated video frames.
|
||||
|
||||
Methods:
|
||||
__init__(self, videos: Union[torch.Tensor, np.ndarray]): Initializes the FaceAnimatePipelineOutput object with the generated video frames.
|
||||
"""
|
||||
videos: Union[torch.Tensor, np.ndarray]
|
||||
|
||||
class FaceAnimatePipeline(DiffusionPipeline):
|
||||
"""
|
||||
FaceAnimatePipeline is a custom DiffusionPipeline for animating faces.
|
||||
|
||||
It inherits from the DiffusionPipeline class and is used to animate faces by
|
||||
utilizing a variational autoencoder (VAE), a reference UNet, a denoising UNet,
|
||||
a face locator, and an image processor. The pipeline is responsible for generating
|
||||
and animating face latents, and decoding the latents to produce the final video output.
|
||||
|
||||
Attributes:
|
||||
vae (VaeImageProcessor): Variational autoencoder for processing images.
|
||||
reference_unet (nn.Module): Reference UNet for mutual self-attention.
|
||||
denoising_unet (nn.Module): Denoising UNet for image denoising.
|
||||
face_locator (nn.Module): Face locator for detecting and cropping faces.
|
||||
image_proj (nn.Module): Image projector for processing images.
|
||||
scheduler (Union[DDIMScheduler, PNDMScheduler, LMSDiscreteScheduler,
|
||||
EulerDiscreteScheduler, EulerAncestralDiscreteScheduler,
|
||||
DPMSolverMultistepScheduler]): Diffusion scheduler for
|
||||
controlling the noise level.
|
||||
|
||||
Methods:
|
||||
__init__(self, vae, reference_unet, denoising_unet, face_locator,
|
||||
image_proj, scheduler): Initializes the FaceAnimatePipeline
|
||||
with the given components and scheduler.
|
||||
prepare_latents(self, batch_size, num_channels_latents, width, height,
|
||||
video_length, dtype, device, generator=None, latents=None):
|
||||
Prepares the initial latents for video generation.
|
||||
prepare_extra_step_kwargs(self, generator, eta): Prepares extra keyword
|
||||
arguments for the scheduler step.
|
||||
decode_latents(self, latents): Decodes the latents to produce the final
|
||||
video output.
|
||||
"""
|
||||
def __init__(
|
||||
self,
|
||||
vae,
|
||||
reference_unet,
|
||||
denoising_unet,
|
||||
scheduler: Union[
|
||||
DDIMScheduler,
|
||||
PNDMScheduler,
|
||||
LMSDiscreteScheduler,
|
||||
EulerDiscreteScheduler,
|
||||
EulerAncestralDiscreteScheduler,
|
||||
DPMSolverMultistepScheduler,
|
||||
],
|
||||
) -> None:
|
||||
super().__init__()
|
||||
|
||||
self.register_modules(
|
||||
vae=vae,
|
||||
reference_unet=reference_unet,
|
||||
denoising_unet=denoising_unet,
|
||||
scheduler=scheduler,
|
||||
)
|
||||
|
||||
self.vae_scale_factor: int = 2 ** (len(self.vae.config.block_out_channels) - 1)
|
||||
|
||||
self.ref_image_processor = VaeImageProcessor(
|
||||
vae_scale_factor=self.vae_scale_factor, do_convert_rgb=True,
|
||||
)
|
||||
|
||||
@property
|
||||
def _execution_device(self):
|
||||
if self.device != torch.device("meta") or not hasattr(self.unet, "_hf_hook"):
|
||||
return self.device
|
||||
for module in self.unet.modules():
|
||||
if (
|
||||
hasattr(module, "_hf_hook")
|
||||
and hasattr(module._hf_hook, "execution_device")
|
||||
and module._hf_hook.execution_device is not None
|
||||
):
|
||||
return torch.device(module._hf_hook.execution_device)
|
||||
return self.device
|
||||
|
||||
def prepare_latents(
|
||||
self,
|
||||
batch_size: int, # Number of videos to generate in parallel
|
||||
num_channels_latents: int, # Number of channels in the latents
|
||||
width: int, # Width of the video frame
|
||||
height: int, # Height of the video frame
|
||||
video_length: int, # Length of the video in frames
|
||||
dtype: torch.dtype, # Data type of the latents
|
||||
device: torch.device, # Device to store the latents on
|
||||
generator: Optional[torch.Generator] = None, # Random number generator for reproducibility
|
||||
latents: Optional[torch.Tensor] = None # Pre-generated latents (optional)
|
||||
):
|
||||
"""
|
||||
Prepares the initial latents for video generation.
|
||||
|
||||
Args:
|
||||
batch_size (int): Number of videos to generate in parallel.
|
||||
num_channels_latents (int): Number of channels in the latents.
|
||||
width (int): Width of the video frame.
|
||||
height (int): Height of the video frame.
|
||||
video_length (int): Length of the video in frames.
|
||||
dtype (torch.dtype): Data type of the latents.
|
||||
device (torch.device): Device to store the latents on.
|
||||
generator (Optional[torch.Generator]): Random number generator for reproducibility.
|
||||
latents (Optional[torch.Tensor]): Pre-generated latents (optional).
|
||||
|
||||
Returns:
|
||||
latents (torch.Tensor): Tensor of shape (batch_size, num_channels_latents, width, height)
|
||||
containing the initial latents for video generation.
|
||||
"""
|
||||
shape = (
|
||||
batch_size,
|
||||
num_channels_latents,
|
||||
video_length,
|
||||
height // self.vae_scale_factor,
|
||||
width // self.vae_scale_factor,
|
||||
)
|
||||
if isinstance(generator, list) and len(generator) != batch_size:
|
||||
raise ValueError(
|
||||
f"You have passed a list of generators of length {len(generator)}, but requested an effective batch"
|
||||
f" size of {batch_size}. Make sure the batch size matches the length of the generators."
|
||||
)
|
||||
|
||||
if latents is None:
|
||||
latents = randn_tensor(
|
||||
shape, generator=generator, device=device, dtype=dtype
|
||||
)
|
||||
else:
|
||||
latents = latents.to(device)
|
||||
|
||||
# scale the initial noise by the standard deviation required by the scheduler
|
||||
latents = latents * self.scheduler.init_noise_sigma
|
||||
return latents
|
||||
|
||||
def prepare_extra_step_kwargs(self, generator, eta):
|
||||
"""
|
||||
Prepares extra keyword arguments for the scheduler step.
|
||||
|
||||
Args:
|
||||
generator (Optional[torch.Generator]): Random number generator for reproducibility.
|
||||
eta (float): The eta (η) parameter used with the DDIMScheduler.
|
||||
It corresponds to η in the DDIM paper (https://arxiv.org/abs/2010.02502) and should be between [0, 1].
|
||||
|
||||
Returns:
|
||||
dict: A dictionary containing the extra keyword arguments for the scheduler step.
|
||||
"""
|
||||
# prepare extra kwargs for the scheduler step, since not all schedulers have the same signature
|
||||
# eta (η) is only used with the DDIMScheduler, it will be ignored for other schedulers.
|
||||
# eta corresponds to η in DDIM paper: https://arxiv.org/abs/2010.02502
|
||||
# and should be between [0, 1]
|
||||
|
||||
accepts_eta = "eta" in set(
|
||||
inspect.signature(self.scheduler.step).parameters.keys()
|
||||
)
|
||||
extra_step_kwargs = {}
|
||||
if accepts_eta:
|
||||
extra_step_kwargs["eta"] = eta
|
||||
|
||||
# check if the scheduler accepts generator
|
||||
accepts_generator = "generator" in set(
|
||||
inspect.signature(self.scheduler.step).parameters.keys()
|
||||
)
|
||||
if accepts_generator:
|
||||
extra_step_kwargs["generator"] = generator
|
||||
return extra_step_kwargs
|
||||
|
||||
def decode_latents(self, latents):
|
||||
"""
|
||||
Decode the latents to produce a video.
|
||||
|
||||
Parameters:
|
||||
latents (torch.Tensor): The latents to be decoded.
|
||||
|
||||
Returns:
|
||||
video (torch.Tensor): The decoded video.
|
||||
video_length (int): The length of the video in frames.
|
||||
"""
|
||||
video_length = latents.shape[2]
|
||||
latents = 1 / 0.18215 * latents
|
||||
latents = rearrange(latents, "b c f h w -> (b f) c h w")
|
||||
# video = self.vae.decode(latents).sample
|
||||
video = []
|
||||
for frame_idx in tqdm(range(latents.shape[0])):
|
||||
video.append(self.vae.decode(
|
||||
latents[frame_idx: frame_idx + 1]).sample)
|
||||
video = torch.cat(video)
|
||||
video = rearrange(video, "(b f) c h w -> b c f h w", f=video_length)
|
||||
video = (video / 2 + 0.5).clamp(0, 1)
|
||||
# we always cast to float32 as this does not cause significant overhead and is compatible with bfloa16
|
||||
video = video.cpu().float().numpy()
|
||||
return video
|
||||
|
||||
|
||||
@torch.no_grad()
|
||||
def __call__(
|
||||
self,
|
||||
gt_encode,
|
||||
audio_tensor,
|
||||
mask_main,
|
||||
width,
|
||||
height,
|
||||
video_length,
|
||||
num_inference_steps,
|
||||
guidance_scale,
|
||||
num_images_per_prompt=1,
|
||||
eta: float = 0.0,
|
||||
motion_scale: Optional[List[torch.Tensor]] = None,
|
||||
generator: Optional[Union[torch.Generator,
|
||||
List[torch.Generator]]] = None,
|
||||
output_type: Optional[str] = "tensor",
|
||||
return_dict: bool = True,
|
||||
callback: Optional[Callable[[
|
||||
int, int, torch.FloatTensor], None]] = None,
|
||||
callback_steps: Optional[int] = 1,
|
||||
**kwargs,
|
||||
):
|
||||
|
||||
# Default height and width to unet
|
||||
height = height or self.unet.config.sample_size * self.vae_scale_factor
|
||||
width = width or self.unet.config.sample_size * self.vae_scale_factor
|
||||
|
||||
device = self._execution_device
|
||||
|
||||
do_classifier_free_guidance = guidance_scale > 1.0
|
||||
|
||||
# Prepare timesteps
|
||||
self.scheduler.set_timesteps(num_inference_steps, device=device)
|
||||
timesteps = self.scheduler.timesteps
|
||||
|
||||
batch_size = 1
|
||||
|
||||
reference_control_writer = ReferenceAttentionControl(
|
||||
self.reference_unet,
|
||||
do_classifier_free_guidance=do_classifier_free_guidance,
|
||||
mode="write",
|
||||
batch_size=batch_size,
|
||||
fusion_blocks="full",
|
||||
)
|
||||
reference_control_reader = ReferenceAttentionControl(
|
||||
self.denoising_unet,
|
||||
do_classifier_free_guidance=do_classifier_free_guidance,
|
||||
mode="read",
|
||||
batch_size=batch_size,
|
||||
fusion_blocks="full",
|
||||
audio_feature_ratio=3.0,
|
||||
)
|
||||
|
||||
num_channels_latents = self.denoising_unet.in_channels
|
||||
|
||||
gt_encode = gt_encode.to(device=device, dtype=mask_main.dtype)
|
||||
|
||||
latents = self.prepare_latents(
|
||||
batch_size * num_images_per_prompt,
|
||||
num_channels_latents,
|
||||
width,
|
||||
height,
|
||||
video_length,
|
||||
gt_encode.dtype,
|
||||
device,
|
||||
generator,
|
||||
)
|
||||
# Prepare extra step kwargs.
|
||||
extra_step_kwargs = self.prepare_extra_step_kwargs(generator, eta)
|
||||
|
||||
gt_encode = self.ref_image_processor.preprocess(
|
||||
gt_encode, height=height, width=width) # (bs, c, width, height)
|
||||
gt_encode = gt_encode.to(dtype=self.vae.dtype, device=self.vae.device) # f*c*h*w
|
||||
|
||||
#do vae for gt image
|
||||
gt_encode_chunks = []
|
||||
f = int(np.ceil(gt_encode.shape[0] / 5))
|
||||
for i in range(f):
|
||||
gt_encode_chunks.append(self.vae.encode(gt_encode[i*5 : (i+1)*5]).latent_dist.mean)
|
||||
gt_encode = torch.cat(gt_encode_chunks, dim=0)
|
||||
|
||||
gt_encode = gt_encode.unsqueeze(0) # b*f*c*h*w
|
||||
gt_encode = gt_encode * 0.18215
|
||||
gt_encode = rearrange(gt_encode, "b f c h w -> b c f h w")
|
||||
if torch.cuda.is_available():
|
||||
torch.cuda.empty_cache()
|
||||
|
||||
#########################################
|
||||
num_warmup_steps = 0
|
||||
context_schedule="uniform"
|
||||
context_frames=12
|
||||
context_stride=1
|
||||
context_overlap=3
|
||||
context_batch_size=1
|
||||
|
||||
context_scheduler = get_context_scheduler(context_schedule)
|
||||
|
||||
#### old
|
||||
# context_queue = list(context_scheduler(0,num_inference_steps,latents.shape[2],context_frames,context_stride,context_overlap,))
|
||||
|
||||
|
||||
|
||||
|
||||
mask_main = mask_main.repeat(1,1,video_length,1,1)
|
||||
|
||||
with self.progress_bar(total=num_inference_steps) as progress_bar:
|
||||
for t_i, t in enumerate(timesteps):
|
||||
|
||||
latents = latents*mask_main + gt_encode*(1-mask_main)
|
||||
|
||||
noise_pred = torch.zeros(
|
||||
(
|
||||
latents.shape[0] * (2 if do_classifier_free_guidance else 1),
|
||||
*latents.shape[1:],
|
||||
),
|
||||
device=latents.device,
|
||||
dtype=latents.dtype,
|
||||
)
|
||||
counter = torch.zeros(
|
||||
(1, 1, latents.shape[2], 1, 1),
|
||||
device=latents.device,
|
||||
dtype=latents.dtype,
|
||||
)
|
||||
#### new
|
||||
offset = t_i % context_frames
|
||||
offset_context_queue = list(context_scheduler(
|
||||
0, num_inference_steps, latents.shape[2] + offset, context_frames, context_stride, context_overlap))
|
||||
|
||||
context_queue = []
|
||||
for ctx_q in offset_context_queue:
|
||||
new_ctx_q = [i_p - offset for i_p in ctx_q]
|
||||
new_ctx_q = [i_p for i_p in new_ctx_q if i_p >= 0 and i_p < latents.shape[2]]
|
||||
context_queue.append(new_ctx_q)
|
||||
|
||||
#########
|
||||
|
||||
num_context_batches = math.ceil(len(context_queue) / context_batch_size)
|
||||
|
||||
global_context = []
|
||||
for j in range(num_context_batches):
|
||||
global_context.append(
|
||||
context_queue[
|
||||
j * context_batch_size : (j + 1) * context_batch_size
|
||||
]
|
||||
)
|
||||
|
||||
for count, context_main in enumerate(global_context):
|
||||
|
||||
context, len_context = modify_context_all(context_main)
|
||||
|
||||
mid_index = int(len(context[0])/2)
|
||||
|
||||
# self.reference_unet(
|
||||
# gt_encode[:,:,context[0][mid_index]].repeat(
|
||||
# context_batch_size * (1), 1, 1, 1),
|
||||
# torch.zeros_like(t),
|
||||
# encoder_hidden_states=None,
|
||||
# return_dict=False,
|
||||
# )
|
||||
|
||||
self.reference_unet(
|
||||
rearrange(gt_encode[0], "c f h w -> f c h w")[context[0]],
|
||||
torch.zeros_like(t),
|
||||
encoder_hidden_states=None,
|
||||
return_dict=False,
|
||||
)
|
||||
|
||||
|
||||
reference_control_reader.update(reference_control_writer, do_classifier_free_guidance=do_classifier_free_guidance)
|
||||
|
||||
latent_model_input = (
|
||||
torch.cat([latents[:, :, c] for c in context])
|
||||
.to(device)
|
||||
.repeat(2 if do_classifier_free_guidance else 1, 1, 1, 1, 1)
|
||||
)
|
||||
c_audio_latents = torch.cat([audio_tensor[:, c] for c in context]).to(device)
|
||||
audio_latents = torch.cat([torch.zeros_like(c_audio_latents), c_audio_latents], 0)
|
||||
|
||||
latent_model_input = self.scheduler.scale_model_input(
|
||||
latent_model_input, t
|
||||
)
|
||||
|
||||
pred = self.denoising_unet(
|
||||
latent_model_input,
|
||||
t,
|
||||
encoder_hidden_states=None,
|
||||
audio_cond_fea=audio_latents if do_classifier_free_guidance else c_audio_latents,
|
||||
return_dict=False,
|
||||
)[0]
|
||||
|
||||
reference_control_writer.clear()
|
||||
reference_control_reader.clear()
|
||||
|
||||
if do_classifier_free_guidance:
|
||||
pred_uc, pred_c = pred.chunk(2)
|
||||
pred = torch.cat([pred_uc.unsqueeze(0), pred_c.unsqueeze(0)])
|
||||
else:
|
||||
pred = pred.unsqueeze(1)
|
||||
|
||||
for j, c in enumerate(context_main):
|
||||
|
||||
if context_main[0][0]==0:
|
||||
noise_pred[:, :, c[:len_context]] = noise_pred[:, :, c[:len_context]] + pred[:, j, :, :len_context]
|
||||
else:
|
||||
noise_pred[:, :, c[:len_context]] = noise_pred[:, :, c[:len_context]] + pred[:, j, :, -len_context:]
|
||||
|
||||
counter[:, :, c[:len_context]] = counter[:, :, c[:len_context]] + 1
|
||||
|
||||
# perform guidance
|
||||
if do_classifier_free_guidance:
|
||||
noise_pred_uncond, noise_pred_text = (noise_pred / counter).chunk(2)
|
||||
noise_pred = noise_pred_uncond + guidance_scale * (
|
||||
noise_pred_text - noise_pred_uncond
|
||||
)
|
||||
|
||||
latents = self.scheduler.step(
|
||||
noise_pred, t, latents, **extra_step_kwargs
|
||||
).prev_sample
|
||||
|
||||
if t_i == len(timesteps) - 1 or (
|
||||
(t_i + 1) > num_warmup_steps and (t_i + 1) % self.scheduler.order == 0
|
||||
):
|
||||
progress_bar.update()
|
||||
|
||||
reference_control_reader.clear()
|
||||
reference_control_writer.clear()
|
||||
|
||||
################################
|
||||
|
||||
latents = latents*mask_main + gt_encode*(1-mask_main)
|
||||
# Post-processing
|
||||
latents = self.decode_latents(latents) # (b, c, f, h, w)
|
||||
|
||||
# Convert to tensor
|
||||
if output_type == "tensor":
|
||||
latents = torch.from_numpy(latents)
|
||||
|
||||
if not return_dict:
|
||||
return latents
|
||||
|
||||
return FaceAnimatePipelineOutput(videos=latents)
|
||||
@@ -0,0 +1,489 @@
|
||||
# pylint: disable=R0801
|
||||
"""
|
||||
This module is responsible for animating faces in videos using a combination of deep learning techniques.
|
||||
It provides a pipeline for generating face animations by processing video frames and extracting face features.
|
||||
The module utilizes various schedulers and utilities for efficient face animation and supports different types
|
||||
of latents for more control over the animation process.
|
||||
|
||||
Functions and Classes:
|
||||
- FaceAnimatePipeline: A class that extends the DiffusionPipeline class from the diffusers library to handle face animation tasks.
|
||||
- __init__: Initializes the pipeline with the necessary components (VAE, UNets, face locator, etc.).
|
||||
- prepare_latents: Generates or loads latents for the animation process, scaling them according to the scheduler's requirements.
|
||||
- prepare_extra_step_kwargs: Prepares extra keyword arguments for the scheduler step, ensuring compatibility with different schedulers.
|
||||
- decode_latents: Decodes the latents into video frames, ready for animation.
|
||||
|
||||
Usage:
|
||||
- Import the necessary packages and classes.
|
||||
- Create a FaceAnimatePipeline instance with the required components.
|
||||
- Prepare the latents for the animation process.
|
||||
- Use the pipeline to generate the animated video.
|
||||
|
||||
Note:
|
||||
- This module is designed to work with the diffusers library, which provides the underlying framework for face animation using deep learning.
|
||||
- The module is intended for research and development purposes, and further optimization and customization may be required for specific use cases.
|
||||
"""
|
||||
|
||||
import inspect
|
||||
from dataclasses import dataclass
|
||||
from typing import Callable, List, Optional, Union
|
||||
import math
|
||||
import numpy as np
|
||||
import torch
|
||||
from diffusers import (DDIMScheduler, DiffusionPipeline,
|
||||
DPMSolverMultistepScheduler,
|
||||
EulerAncestralDiscreteScheduler, EulerDiscreteScheduler,
|
||||
LMSDiscreteScheduler, PNDMScheduler)
|
||||
from diffusers.image_processor import VaeImageProcessor
|
||||
from diffusers.utils import BaseOutput
|
||||
from diffusers.utils.torch_utils import randn_tensor
|
||||
from einops import rearrange, repeat
|
||||
from tqdm import tqdm
|
||||
|
||||
from .context import get_context_scheduler, modify_context_all
|
||||
from ..models.mutual_self_attention import ReferenceAttentionControl
|
||||
|
||||
|
||||
@dataclass
|
||||
class FaceAnimatePipelineOutput(BaseOutput):
|
||||
"""
|
||||
FaceAnimatePipelineOutput is a custom class that inherits from BaseOutput and represents the output of the FaceAnimatePipeline.
|
||||
|
||||
Attributes:
|
||||
videos (Union[torch.Tensor, np.ndarray]): A tensor or numpy array containing the generated video frames.
|
||||
|
||||
Methods:
|
||||
__init__(self, videos: Union[torch.Tensor, np.ndarray]): Initializes the FaceAnimatePipelineOutput object with the generated video frames.
|
||||
"""
|
||||
videos: Union[torch.Tensor, np.ndarray]
|
||||
|
||||
class FaceAnimatePipeline(DiffusionPipeline):
|
||||
"""
|
||||
FaceAnimatePipeline is a custom DiffusionPipeline for animating faces.
|
||||
|
||||
It inherits from the DiffusionPipeline class and is used to animate faces by
|
||||
utilizing a variational autoencoder (VAE), a reference UNet, a denoising UNet,
|
||||
a face locator, and an image processor. The pipeline is responsible for generating
|
||||
and animating face latents, and decoding the latents to produce the final video output.
|
||||
|
||||
Attributes:
|
||||
vae (VaeImageProcessor): Variational autoencoder for processing images.
|
||||
reference_unet (nn.Module): Reference UNet for mutual self-attention.
|
||||
denoising_unet (nn.Module): Denoising UNet for image denoising.
|
||||
face_locator (nn.Module): Face locator for detecting and cropping faces.
|
||||
image_proj (nn.Module): Image projector for processing images.
|
||||
scheduler (Union[DDIMScheduler, PNDMScheduler, LMSDiscreteScheduler,
|
||||
EulerDiscreteScheduler, EulerAncestralDiscreteScheduler,
|
||||
DPMSolverMultistepScheduler]): Diffusion scheduler for
|
||||
controlling the noise level.
|
||||
|
||||
Methods:
|
||||
__init__(self, vae, reference_unet, denoising_unet, face_locator,
|
||||
image_proj, scheduler): Initializes the FaceAnimatePipeline
|
||||
with the given components and scheduler.
|
||||
prepare_latents(self, batch_size, num_channels_latents, width, height,
|
||||
video_length, dtype, device, generator=None, latents=None):
|
||||
Prepares the initial latents for video generation.
|
||||
prepare_extra_step_kwargs(self, generator, eta): Prepares extra keyword
|
||||
arguments for the scheduler step.
|
||||
decode_latents(self, latents): Decodes the latents to produce the final
|
||||
video output.
|
||||
"""
|
||||
def __init__(
|
||||
self,
|
||||
vae,
|
||||
reference_unet,
|
||||
denoising_unet,
|
||||
scheduler: Union[
|
||||
DDIMScheduler,
|
||||
PNDMScheduler,
|
||||
LMSDiscreteScheduler,
|
||||
EulerDiscreteScheduler,
|
||||
EulerAncestralDiscreteScheduler,
|
||||
DPMSolverMultistepScheduler,
|
||||
]
|
||||
) -> None:
|
||||
super().__init__()
|
||||
|
||||
self.register_modules(
|
||||
vae=vae,
|
||||
reference_unet=reference_unet,
|
||||
denoising_unet=denoising_unet,
|
||||
scheduler=scheduler,
|
||||
)
|
||||
|
||||
self.vae_scale_factor: int = 2 ** (len(self.vae.config.block_out_channels) - 1)
|
||||
|
||||
self.ref_image_processor = VaeImageProcessor(
|
||||
vae_scale_factor=self.vae_scale_factor, do_convert_rgb=True,
|
||||
)
|
||||
|
||||
self.context_connection_frames = torch.zeros((20,4,3,64,64))
|
||||
self.context_connection_frames_prev = None
|
||||
|
||||
@property
|
||||
def _execution_device(self):
|
||||
if self.device != torch.device("meta") or not hasattr(self.unet, "_hf_hook"):
|
||||
return self.device
|
||||
for module in self.unet.modules():
|
||||
if (
|
||||
hasattr(module, "_hf_hook")
|
||||
and hasattr(module._hf_hook, "execution_device")
|
||||
and module._hf_hook.execution_device is not None
|
||||
):
|
||||
return torch.device(module._hf_hook.execution_device)
|
||||
return self.device
|
||||
|
||||
def prepare_latents(
|
||||
self,
|
||||
batch_size: int, # Number of videos to generate in parallel
|
||||
num_channels_latents: int, # Number of channels in the latents
|
||||
width: int, # Width of the video frame
|
||||
height: int, # Height of the video frame
|
||||
video_length: int, # Length of the video in frames
|
||||
dtype: torch.dtype, # Data type of the latents
|
||||
device: torch.device, # Device to store the latents on
|
||||
generator: Optional[torch.Generator] = None, # Random number generator for reproducibility
|
||||
latents: Optional[torch.Tensor] = None # Pre-generated latents (optional)
|
||||
):
|
||||
"""
|
||||
Prepares the initial latents for video generation.
|
||||
|
||||
Args:
|
||||
batch_size (int): Number of videos to generate in parallel.
|
||||
num_channels_latents (int): Number of channels in the latents.
|
||||
width (int): Width of the video frame.
|
||||
height (int): Height of the video frame.
|
||||
video_length (int): Length of the video in frames.
|
||||
dtype (torch.dtype): Data type of the latents.
|
||||
device (torch.device): Device to store the latents on.
|
||||
generator (Optional[torch.Generator]): Random number generator for reproducibility.
|
||||
latents (Optional[torch.Tensor]): Pre-generated latents (optional).
|
||||
|
||||
Returns:
|
||||
latents (torch.Tensor): Tensor of shape (batch_size, num_channels_latents, width, height)
|
||||
containing the initial latents for video generation.
|
||||
"""
|
||||
shape = (
|
||||
batch_size,
|
||||
num_channels_latents,
|
||||
video_length,
|
||||
height // self.vae_scale_factor,
|
||||
width // self.vae_scale_factor,
|
||||
)
|
||||
if isinstance(generator, list) and len(generator) != batch_size:
|
||||
raise ValueError(
|
||||
f"You have passed a list of generators of length {len(generator)}, but requested an effective batch"
|
||||
f" size of {batch_size}. Make sure the batch size matches the length of the generators."
|
||||
)
|
||||
|
||||
if latents is None:
|
||||
latents = randn_tensor(
|
||||
shape, generator=generator, device=device, dtype=dtype
|
||||
)
|
||||
else:
|
||||
latents = latents.to(device)
|
||||
|
||||
# scale the initial noise by the standard deviation required by the scheduler
|
||||
latents = latents * self.scheduler.init_noise_sigma
|
||||
return latents
|
||||
|
||||
def prepare_extra_step_kwargs(self, generator, eta):
|
||||
"""
|
||||
Prepares extra keyword arguments for the scheduler step.
|
||||
|
||||
Args:
|
||||
generator (Optional[torch.Generator]): Random number generator for reproducibility.
|
||||
eta (float): The eta (η) parameter used with the DDIMScheduler.
|
||||
It corresponds to η in the DDIM paper (https://arxiv.org/abs/2010.02502) and should be between [0, 1].
|
||||
|
||||
Returns:
|
||||
dict: A dictionary containing the extra keyword arguments for the scheduler step.
|
||||
"""
|
||||
# prepare extra kwargs for the scheduler step, since not all schedulers have the same signature
|
||||
# eta (η) is only used with the DDIMScheduler, it will be ignored for other schedulers.
|
||||
# eta corresponds to η in DDIM paper: https://arxiv.org/abs/2010.02502
|
||||
# and should be between [0, 1]
|
||||
|
||||
accepts_eta = "eta" in set(
|
||||
inspect.signature(self.scheduler.step).parameters.keys()
|
||||
)
|
||||
extra_step_kwargs = {}
|
||||
if accepts_eta:
|
||||
extra_step_kwargs["eta"] = eta
|
||||
|
||||
# check if the scheduler accepts generator
|
||||
accepts_generator = "generator" in set(
|
||||
inspect.signature(self.scheduler.step).parameters.keys()
|
||||
)
|
||||
if accepts_generator:
|
||||
extra_step_kwargs["generator"] = generator
|
||||
return extra_step_kwargs
|
||||
|
||||
def decode_latents(self, latents):
|
||||
"""
|
||||
Decode the latents to produce a video.
|
||||
|
||||
Parameters:
|
||||
latents (torch.Tensor): The latents to be decoded.
|
||||
|
||||
Returns:
|
||||
video (torch.Tensor): The decoded video.
|
||||
video_length (int): The length of the video in frames.
|
||||
"""
|
||||
video_length = latents.shape[2]
|
||||
latents = 1 / 0.18215 * latents
|
||||
latents = rearrange(latents, "b c f h w -> (b f) c h w")
|
||||
# video = self.vae.decode(latents).sample
|
||||
video = []
|
||||
for frame_idx in tqdm(range(latents.shape[0])):
|
||||
video.append(self.vae.decode(
|
||||
latents[frame_idx: frame_idx + 1]).sample)
|
||||
video = torch.cat(video)
|
||||
video = rearrange(video, "(b f) c h w -> b c f h w", f=video_length)
|
||||
video = (video / 2 + 0.5).clamp(0, 1)
|
||||
# we always cast to float32 as this does not cause significant overhead and is compatible with bfloa16
|
||||
video = video.cpu().float().numpy()
|
||||
return video
|
||||
|
||||
|
||||
@torch.no_grad()
|
||||
def __call__(
|
||||
self,
|
||||
gt_encode,
|
||||
audio_tensor,
|
||||
connection,
|
||||
mask_main,
|
||||
width,
|
||||
height,
|
||||
video_length,
|
||||
num_inference_steps,
|
||||
guidance_scale,
|
||||
num_images_per_prompt=1,
|
||||
eta: float = 0.0,
|
||||
motion_scale: Optional[List[torch.Tensor]] = None,
|
||||
generator: Optional[Union[torch.Generator,
|
||||
List[torch.Generator]]] = None,
|
||||
output_type: Optional[str] = "tensor",
|
||||
return_dict: bool = True,
|
||||
callback: Optional[Callable[[
|
||||
int, int, torch.FloatTensor], None]] = None,
|
||||
callback_steps: Optional[int] = 1,
|
||||
**kwargs,
|
||||
):
|
||||
# Default height and width to unet
|
||||
height = height or self.unet.config.sample_size * self.vae_scale_factor
|
||||
width = width or self.unet.config.sample_size * self.vae_scale_factor
|
||||
|
||||
device = self._execution_device
|
||||
|
||||
do_classifier_free_guidance = guidance_scale > 1.0
|
||||
|
||||
# Prepare timesteps
|
||||
self.scheduler.set_timesteps(num_inference_steps, device=device)
|
||||
timesteps = self.scheduler.timesteps
|
||||
|
||||
batch_size = 1
|
||||
|
||||
reference_control_writer = ReferenceAttentionControl(
|
||||
self.reference_unet,
|
||||
do_classifier_free_guidance=do_classifier_free_guidance,
|
||||
mode="write",
|
||||
batch_size=batch_size,
|
||||
fusion_blocks="full",
|
||||
)
|
||||
reference_control_reader = ReferenceAttentionControl(
|
||||
self.denoising_unet,
|
||||
do_classifier_free_guidance=do_classifier_free_guidance,
|
||||
mode="read",
|
||||
batch_size=batch_size,
|
||||
fusion_blocks="full",
|
||||
audio_feature_ratio=3.0,
|
||||
)
|
||||
|
||||
num_channels_latents = self.denoising_unet.in_channels
|
||||
|
||||
gt_encode = gt_encode.to(device=device, dtype=mask_main.dtype)
|
||||
|
||||
latents = self.prepare_latents(
|
||||
batch_size * num_images_per_prompt,
|
||||
num_channels_latents,
|
||||
width,
|
||||
height,
|
||||
video_length,
|
||||
gt_encode.dtype,
|
||||
device,
|
||||
generator,
|
||||
)
|
||||
|
||||
self.context_connection_frames[0] = latents[0,:,-3:]
|
||||
|
||||
# Prepare extra step kwargs.
|
||||
extra_step_kwargs = self.prepare_extra_step_kwargs(generator, eta)
|
||||
|
||||
gt_encode = self.ref_image_processor.preprocess(
|
||||
gt_encode, height=height, width=width) # (bs, c, width, height)
|
||||
gt_encode = gt_encode.to(dtype=self.vae.dtype, device=self.vae.device) # f*c*h*w
|
||||
#do vae for gt image
|
||||
#do vae for gt image
|
||||
gt_encode_chunks = []
|
||||
f = int(np.ceil(gt_encode.shape[0] / 5))
|
||||
for i in range(f):
|
||||
gt_encode_chunks.append(self.vae.encode(gt_encode[i*5 : (i+1)*5]).latent_dist.mean)
|
||||
gt_encode = torch.cat(gt_encode_chunks, dim=0)
|
||||
|
||||
gt_encode = gt_encode.unsqueeze(0) # b*c*f*h*w
|
||||
gt_encode = gt_encode * 0.18215
|
||||
gt_encode = rearrange(gt_encode, "b f c h w -> b c f h w")
|
||||
|
||||
#########################################
|
||||
num_warmup_steps = 0
|
||||
context_schedule="uniform"
|
||||
context_frames=12
|
||||
context_stride=1
|
||||
context_overlap=3
|
||||
context_batch_size=1
|
||||
|
||||
context_scheduler = get_context_scheduler(context_schedule)
|
||||
|
||||
#### old
|
||||
# context_queue = list(context_scheduler(0,num_inference_steps,latents.shape[2],context_frames,context_stride,context_overlap,))
|
||||
|
||||
mask_main = mask_main.repeat(1,1,video_length,1,1)
|
||||
|
||||
with self.progress_bar(total=num_inference_steps) as progress_bar:
|
||||
for t_i, t in enumerate(timesteps):
|
||||
|
||||
if connection:
|
||||
latents[0,:,:3] = self.context_connection_frames_prev[t_i]
|
||||
|
||||
latents = latents*mask_main + gt_encode*(1-mask_main)
|
||||
|
||||
noise_pred = torch.zeros(
|
||||
(
|
||||
latents.shape[0] * (2 if do_classifier_free_guidance else 1),
|
||||
*latents.shape[1:],
|
||||
),
|
||||
device=latents.device,
|
||||
dtype=latents.dtype,
|
||||
)
|
||||
counter = torch.zeros(
|
||||
(1, 1, latents.shape[2], 1, 1),
|
||||
device=latents.device,
|
||||
dtype=latents.dtype,
|
||||
)
|
||||
|
||||
#### new
|
||||
offset = t_i % context_frames
|
||||
offset_context_queue = list(context_scheduler(
|
||||
0, num_inference_steps, latents.shape[2] + offset, context_frames, context_stride, context_overlap))
|
||||
|
||||
context_queue = []
|
||||
for ctx_q in offset_context_queue:
|
||||
new_ctx_q = [i_p - offset for i_p in ctx_q]
|
||||
new_ctx_q = [i_p for i_p in new_ctx_q if i_p >= 0 and i_p < latents.shape[2]]
|
||||
context_queue.append(new_ctx_q)
|
||||
|
||||
#########
|
||||
|
||||
num_context_batches = math.ceil(len(context_queue) / context_batch_size)
|
||||
|
||||
global_context = []
|
||||
for j in range(num_context_batches):
|
||||
global_context.append(
|
||||
context_queue[
|
||||
j * context_batch_size : (j + 1) * context_batch_size
|
||||
]
|
||||
)
|
||||
|
||||
for count, context_main in enumerate(global_context):
|
||||
|
||||
context, len_context = modify_context_all(context_main)
|
||||
|
||||
mid_index = int(len(context[0])/2)
|
||||
|
||||
self.reference_unet(
|
||||
gt_encode[:,:,context[0][mid_index]].repeat(
|
||||
context_batch_size * (1), 1, 1, 1),
|
||||
torch.zeros_like(t),
|
||||
encoder_hidden_states=None,
|
||||
return_dict=False,
|
||||
)
|
||||
reference_control_reader.update(reference_control_writer, do_classifier_free_guidance=do_classifier_free_guidance)
|
||||
|
||||
latent_model_input = (
|
||||
torch.cat([latents[:, :, c] for c in context])
|
||||
.to(device)
|
||||
.repeat(2 if do_classifier_free_guidance else 1, 1, 1, 1, 1)
|
||||
)
|
||||
c_audio_latents = torch.cat([audio_tensor[:, c] for c in context]).to(device)
|
||||
audio_latents = torch.cat([torch.zeros_like(c_audio_latents), c_audio_latents], 0)
|
||||
|
||||
latent_model_input = self.scheduler.scale_model_input(
|
||||
latent_model_input, t
|
||||
)
|
||||
|
||||
pred = self.denoising_unet(
|
||||
latent_model_input,
|
||||
t,
|
||||
encoder_hidden_states=None,
|
||||
audio_cond_fea=audio_latents if do_classifier_free_guidance else c_audio_latents,
|
||||
return_dict=False,
|
||||
)[0]
|
||||
|
||||
reference_control_writer.clear()
|
||||
reference_control_reader.clear()
|
||||
|
||||
if do_classifier_free_guidance:
|
||||
pred_uc, pred_c = pred.chunk(2)
|
||||
pred = torch.cat([pred_uc.unsqueeze(0), pred_c.unsqueeze(0)])
|
||||
else:
|
||||
pred = pred.unsqueeze(1)
|
||||
|
||||
for j, c in enumerate(context_main):
|
||||
|
||||
if context_main[0][0]==0:
|
||||
noise_pred[:, :, c[:len_context]] = noise_pred[:, :, c[:len_context]] + pred[:, j, :, :len_context]
|
||||
else:
|
||||
noise_pred[:, :, c[:len_context]] = noise_pred[:, :, c[:len_context]] + pred[:, j, :, -len_context:]
|
||||
|
||||
counter[:, :, c[:len_context]] = counter[:, :, c[:len_context]] + 1
|
||||
|
||||
# perform guidance
|
||||
if do_classifier_free_guidance:
|
||||
noise_pred_uncond, noise_pred_text = (noise_pred / counter).chunk(2)
|
||||
noise_pred = noise_pred_uncond + guidance_scale * (
|
||||
noise_pred_text - noise_pred_uncond
|
||||
)
|
||||
|
||||
latents = self.scheduler.step(
|
||||
noise_pred, t, latents, **extra_step_kwargs
|
||||
).prev_sample
|
||||
|
||||
if t_i == len(timesteps) - 1 or (
|
||||
(t_i + 1) > num_warmup_steps and (t_i + 1) % self.scheduler.order == 0
|
||||
):
|
||||
progress_bar.update()
|
||||
|
||||
if t_i < (len(timesteps) - 1):
|
||||
self.context_connection_frames[t_i+1] = latents[0,:,-3:]
|
||||
|
||||
reference_control_reader.clear()
|
||||
reference_control_writer.clear()
|
||||
|
||||
################################
|
||||
|
||||
self.context_connection_frames_prev = self.context_connection_frames
|
||||
self.context_connection_frames = torch.zeros((num_inference_steps,4,3,64,64))
|
||||
|
||||
latents = latents*mask_main + gt_encode*(1-mask_main)
|
||||
# Post-processing
|
||||
latents = self.decode_latents(latents) # (b, c, f, h, w)
|
||||
|
||||
# Convert to tensor
|
||||
if output_type == "tensor":
|
||||
latents = torch.from_numpy(latents)
|
||||
|
||||
if not return_dict:
|
||||
return latents
|
||||
|
||||
return FaceAnimatePipelineOutput(videos=latents)
|
||||
@@ -0,0 +1,476 @@
|
||||
# pylint: disable=R0801
|
||||
"""
|
||||
This module is responsible for handling the animation of faces using a combination of deep learning models and image processing techniques.
|
||||
It provides a pipeline to generate realistic face animations by incorporating user-provided conditions such as facial expressions and environments.
|
||||
The module utilizes various schedulers and utilities to optimize the animation process and ensure efficient performance.
|
||||
|
||||
Functions and Classes:
|
||||
- StaticPipelineOutput: A class that represents the output of the animation pipeline, c
|
||||
ontaining properties and methods related to the generated images.
|
||||
- prepare_latents: A function that prepares the initial noise for the animation process,
|
||||
scaling it according to the scheduler's requirements.
|
||||
- prepare_condition: A function that processes the user-provided conditions
|
||||
(e.g., facial expressions) and prepares them for use in the animation pipeline.
|
||||
- decode_latents: A function that decodes the latent representations of the face animations into
|
||||
their corresponding image formats.
|
||||
- prepare_extra_step_kwargs: A function that prepares additional parameters for each step of
|
||||
the animation process, such as the generator and eta values.
|
||||
|
||||
Dependencies:
|
||||
- numpy: A library for numerical computing.
|
||||
- torch: A machine learning library based on PyTorch.
|
||||
- diffusers: A library for image-to-image diffusion models.
|
||||
- transformers: A library for pre-trained transformer models.
|
||||
|
||||
Usage:
|
||||
- To create an instance of the animation pipeline, provide the necessary components such as
|
||||
the VAE, reference UNET, denoising UNET, face locator, and image processor.
|
||||
- Use the pipeline's methods to prepare the latents, conditions, and extra step arguments as
|
||||
required for the animation process.
|
||||
- Generate the face animations by decoding the latents and processing the conditions.
|
||||
|
||||
Note:
|
||||
- The module is designed to work with the diffusers library, which is based on
|
||||
the paper "Diffusion Models for Image-to-Image Translation" (https://arxiv.org/abs/2102.02765).
|
||||
- The face animations generated by this module should be used for entertainment purposes
|
||||
only and should respect the rights and privacy of the individuals involved.
|
||||
"""
|
||||
import inspect
|
||||
from dataclasses import dataclass
|
||||
from typing import Callable, List, Optional, Union
|
||||
|
||||
import numpy as np
|
||||
import torch
|
||||
from diffusers import DiffusionPipeline
|
||||
from diffusers.image_processor import VaeImageProcessor
|
||||
from diffusers.schedulers import (DDIMScheduler, DPMSolverMultistepScheduler,
|
||||
EulerAncestralDiscreteScheduler,
|
||||
EulerDiscreteScheduler, LMSDiscreteScheduler,
|
||||
PNDMScheduler)
|
||||
from diffusers.utils import BaseOutput, is_accelerate_available
|
||||
from diffusers.utils.torch_utils import randn_tensor
|
||||
from einops import rearrange
|
||||
from tqdm import tqdm
|
||||
from transformers import CLIPImageProcessor
|
||||
|
||||
from ..models.mutual_self_attention import ReferenceAttentionControl
|
||||
|
||||
if is_accelerate_available():
|
||||
from accelerate import cpu_offload
|
||||
else:
|
||||
raise ImportError("Please install accelerate via `pip install accelerate`")
|
||||
|
||||
|
||||
@dataclass
|
||||
class StaticPipelineOutput(BaseOutput):
|
||||
"""
|
||||
StaticPipelineOutput is a class that represents the output of the static pipeline.
|
||||
It contains the images generated by the pipeline as a union of torch.Tensor and np.ndarray.
|
||||
|
||||
Attributes:
|
||||
images (Union[torch.Tensor, np.ndarray]): The generated images.
|
||||
"""
|
||||
images: Union[torch.Tensor, np.ndarray]
|
||||
|
||||
|
||||
class StaticPipeline(DiffusionPipeline):
|
||||
"""
|
||||
StaticPipelineOutput is a class that represents the output of the static pipeline.
|
||||
It contains the images generated by the pipeline as a union of torch.Tensor and np.ndarray.
|
||||
|
||||
Attributes:
|
||||
images (Union[torch.Tensor, np.ndarray]): The generated images.
|
||||
"""
|
||||
_optional_components = []
|
||||
|
||||
def __init__(
|
||||
self,
|
||||
vae,
|
||||
reference_unet,
|
||||
denoising_unet,
|
||||
scheduler: Union[
|
||||
DDIMScheduler,
|
||||
PNDMScheduler,
|
||||
LMSDiscreteScheduler,
|
||||
EulerDiscreteScheduler,
|
||||
EulerAncestralDiscreteScheduler,
|
||||
DPMSolverMultistepScheduler,
|
||||
],
|
||||
):
|
||||
super().__init__()
|
||||
|
||||
self.register_modules(
|
||||
vae=vae,
|
||||
reference_unet=reference_unet,
|
||||
denoising_unet=denoising_unet,
|
||||
scheduler=scheduler,
|
||||
)
|
||||
self.vae_scale_factor = 2 ** (
|
||||
len(self.vae.config.block_out_channels) - 1)
|
||||
self.clip_image_processor = CLIPImageProcessor()
|
||||
self.ref_image_processor = VaeImageProcessor(
|
||||
vae_scale_factor=self.vae_scale_factor, do_convert_rgb=True
|
||||
)
|
||||
self.cond_image_processor = VaeImageProcessor(
|
||||
vae_scale_factor=self.vae_scale_factor,
|
||||
do_convert_rgb=True,
|
||||
do_normalize=False,
|
||||
)
|
||||
|
||||
def enable_vae_slicing(self):
|
||||
"""
|
||||
Enable VAE slicing.
|
||||
|
||||
This method enables slicing for the VAE model, which can help improve the performance of decoding latents when working with large images.
|
||||
"""
|
||||
self.vae.enable_slicing()
|
||||
|
||||
def disable_vae_slicing(self):
|
||||
"""
|
||||
Disable vae slicing.
|
||||
|
||||
This function disables the vae slicing for the StaticPipeline object.
|
||||
It calls the `disable_slicing()` method of the vae model.
|
||||
This is useful when you want to use the entire vae model for decoding latents
|
||||
instead of slicing it for better performance.
|
||||
"""
|
||||
self.vae.disable_slicing()
|
||||
|
||||
def enable_sequential_cpu_offload(self, gpu_id=0):
|
||||
"""
|
||||
Offloads selected models to the GPU for increased performance.
|
||||
|
||||
Args:
|
||||
gpu_id (int, optional): The ID of the GPU to offload models to. Defaults to 0.
|
||||
"""
|
||||
device = torch.device(f"cuda:{gpu_id}")
|
||||
|
||||
for cpu_offloaded_model in [self.unet, self.text_encoder, self.vae]:
|
||||
if cpu_offloaded_model is not None:
|
||||
cpu_offload(cpu_offloaded_model, device)
|
||||
|
||||
@property
|
||||
def _execution_device(self):
|
||||
if self.device != torch.device("meta") or not hasattr(self.unet, "_hf_hook"):
|
||||
return self.device
|
||||
for module in self.unet.modules():
|
||||
if (
|
||||
hasattr(module, "_hf_hook")
|
||||
and hasattr(module._hf_hook, "execution_device")
|
||||
and module._hf_hook.execution_device is not None
|
||||
):
|
||||
return torch.device(module._hf_hook.execution_device)
|
||||
return self.device
|
||||
|
||||
def decode_latents(self, latents):
|
||||
"""
|
||||
Decode the given latents to video frames.
|
||||
|
||||
Parameters:
|
||||
latents (torch.Tensor): The latents to be decoded. Shape: (batch_size, num_channels_latents, video_length, height, width).
|
||||
|
||||
Returns:
|
||||
video (torch.Tensor): The decoded video frames. Shape: (batch_size, num_channels_latents, video_length, height, width).
|
||||
"""
|
||||
video_length = latents.shape[2]
|
||||
latents = 1 / 0.18215 * latents
|
||||
latents = rearrange(latents, "b c f h w -> (b f) c h w")
|
||||
# video = self.vae.decode(latents).sample
|
||||
video = []
|
||||
for frame_idx in tqdm(range(latents.shape[0])):
|
||||
video.append(self.vae.decode(
|
||||
latents[frame_idx: frame_idx + 1]).sample)
|
||||
video = torch.cat(video)
|
||||
video = rearrange(video, "(b f) c h w -> b c f h w", f=video_length)
|
||||
video = (video / 2 + 0.5).clamp(0, 1)
|
||||
# we always cast to float32 as this does not cause significant overhead and is compatible with bfloa16
|
||||
video = video.cpu().float().numpy()
|
||||
return video
|
||||
|
||||
def prepare_extra_step_kwargs(self, generator, eta):
|
||||
"""
|
||||
Prepare extra keyword arguments for the scheduler step.
|
||||
|
||||
Since not all schedulers have the same signature, this function helps to create a consistent interface for the scheduler.
|
||||
|
||||
Args:
|
||||
generator (Optional[torch.Generator]): A random number generator for reproducibility.
|
||||
eta (float): The eta parameter used with the DDIMScheduler. It should be between 0 and 1.
|
||||
|
||||
Returns:
|
||||
dict: A dictionary containing the extra keyword arguments for the scheduler step.
|
||||
"""
|
||||
# eta (η) is only used with the DDIMScheduler, it will be ignored for other schedulers.
|
||||
# eta corresponds to η in DDIM paper: https://arxiv.org/abs/2010.02502
|
||||
# and should be between [0, 1]
|
||||
|
||||
accepts_eta = "eta" in set(
|
||||
inspect.signature(self.scheduler.step).parameters.keys()
|
||||
)
|
||||
extra_step_kwargs = {}
|
||||
if accepts_eta:
|
||||
extra_step_kwargs["eta"] = eta
|
||||
|
||||
# check if the scheduler accepts generator
|
||||
accepts_generator = "generator" in set(
|
||||
inspect.signature(self.scheduler.step).parameters.keys()
|
||||
)
|
||||
if accepts_generator:
|
||||
extra_step_kwargs["generator"] = generator
|
||||
return extra_step_kwargs
|
||||
|
||||
def prepare_latents(
|
||||
self,
|
||||
batch_size,
|
||||
num_channels_latents,
|
||||
width,
|
||||
height,
|
||||
dtype,
|
||||
device,
|
||||
generator,
|
||||
latents=None,
|
||||
):
|
||||
"""
|
||||
Prepares the initial latents for the diffusion pipeline.
|
||||
|
||||
Args:
|
||||
batch_size (int): The number of images to generate in one forward pass.
|
||||
num_channels_latents (int): The number of channels in the latents tensor.
|
||||
width (int): The width of the latents tensor.
|
||||
height (int): The height of the latents tensor.
|
||||
dtype (torch.dtype): The data type of the latents tensor.
|
||||
device (torch.device): The device to place the latents tensor on.
|
||||
generator (Optional[torch.Generator], optional): A random number generator
|
||||
for reproducibility. Defaults to None.
|
||||
latents (Optional[torch.Tensor], optional): Pre-computed latents to use as
|
||||
initial conditions for the diffusion process. Defaults to None.
|
||||
|
||||
Returns:
|
||||
torch.Tensor: The prepared latents tensor.
|
||||
"""
|
||||
shape = (
|
||||
batch_size,
|
||||
num_channels_latents,
|
||||
height // self.vae_scale_factor,
|
||||
width // self.vae_scale_factor,
|
||||
)
|
||||
if isinstance(generator, list) and len(generator) != batch_size:
|
||||
raise ValueError(
|
||||
f"You have passed a list of generators of length {len(generator)}, but requested an effective batch"
|
||||
f" size of {batch_size}. Make sure the batch size matches the length of the generators."
|
||||
)
|
||||
|
||||
if latents is None:
|
||||
latents = randn_tensor(
|
||||
shape, generator=generator, device=device, dtype=dtype
|
||||
)
|
||||
else:
|
||||
latents = latents.to(device)
|
||||
|
||||
# scale the initial noise by the standard deviation required by the scheduler
|
||||
latents = latents * self.scheduler.init_noise_sigma
|
||||
return latents
|
||||
|
||||
def prepare_condition(
|
||||
self,
|
||||
cond_image,
|
||||
width,
|
||||
height,
|
||||
device,
|
||||
dtype,
|
||||
do_classififer_free_guidance=False,
|
||||
):
|
||||
"""
|
||||
Prepares the condition for the face animation pipeline.
|
||||
|
||||
Args:
|
||||
cond_image (torch.Tensor): The conditional image tensor.
|
||||
width (int): The width of the output image.
|
||||
height (int): The height of the output image.
|
||||
device (torch.device): The device to run the pipeline on.
|
||||
dtype (torch.dtype): The data type of the tensor.
|
||||
do_classififer_free_guidance (bool, optional): Whether to use classifier-free guidance or not. Defaults to False.
|
||||
|
||||
Returns:
|
||||
Tuple[torch.Tensor, torch.Tensor]: A tuple of processed condition and mask tensors.
|
||||
"""
|
||||
image = self.cond_image_processor.preprocess(
|
||||
cond_image, height=height, width=width
|
||||
).to(dtype=torch.float32)
|
||||
|
||||
image = image.to(device=device, dtype=dtype)
|
||||
|
||||
if do_classififer_free_guidance:
|
||||
image = torch.cat([image] * 2)
|
||||
|
||||
return image
|
||||
|
||||
@torch.no_grad()
|
||||
def __call__(
|
||||
self,
|
||||
ref_image,
|
||||
gt_image,
|
||||
audio_tensor,
|
||||
mask,
|
||||
width,
|
||||
height,
|
||||
num_inference_steps,
|
||||
guidance_scale,
|
||||
num_images_per_prompt=1,
|
||||
eta: float = 0.0,
|
||||
generator: Optional[Union[torch.Generator,
|
||||
List[torch.Generator]]] = None,
|
||||
output_type: Optional[str] = "tensor",
|
||||
return_dict: bool = True,
|
||||
callback: Optional[Callable[[
|
||||
int, int, torch.FloatTensor], None]] = None,
|
||||
callback_steps: Optional[int] = 1,
|
||||
**kwargs,
|
||||
):
|
||||
# Default height and width to unet
|
||||
height = height or self.unet.config.sample_size * self.vae_scale_factor
|
||||
width = width or self.unet.config.sample_size * self.vae_scale_factor
|
||||
|
||||
device = self._execution_device
|
||||
|
||||
do_classifier_free_guidance = guidance_scale > 1.0
|
||||
|
||||
# Prepare timesteps
|
||||
self.scheduler.set_timesteps(num_inference_steps, device=device)
|
||||
timesteps = self.scheduler.timesteps
|
||||
|
||||
batch_size = 1
|
||||
|
||||
reference_control_writer = ReferenceAttentionControl(
|
||||
self.reference_unet,
|
||||
do_classifier_free_guidance=do_classifier_free_guidance,
|
||||
mode="write",
|
||||
batch_size=batch_size,
|
||||
fusion_blocks="full",
|
||||
)
|
||||
reference_control_reader = ReferenceAttentionControl(
|
||||
self.denoising_unet,
|
||||
do_classifier_free_guidance=do_classifier_free_guidance,
|
||||
mode="read",
|
||||
batch_size=batch_size,
|
||||
fusion_blocks="full",
|
||||
audio_feature_ratio=3.0,
|
||||
)
|
||||
|
||||
num_channels_latents = self.denoising_unet.in_channels
|
||||
latents = self.prepare_latents(
|
||||
batch_size * num_images_per_prompt,
|
||||
num_channels_latents,
|
||||
width,
|
||||
height,
|
||||
self.vae.dtype,
|
||||
device,
|
||||
generator,
|
||||
)
|
||||
latents = latents.unsqueeze(2) # (bs, c, 1, h', w')
|
||||
# latents_dtype = latents.dtype
|
||||
|
||||
# Prepare extra step kwargs.
|
||||
extra_step_kwargs = self.prepare_extra_step_kwargs(generator, eta)
|
||||
|
||||
# Prepare ref image latents
|
||||
# ref_image_tensor = ref_image
|
||||
ref_image_tensor = self.ref_image_processor.preprocess(
|
||||
ref_image, height=height, width=width
|
||||
) # (bs, c, width, height)
|
||||
ref_image_tensor = ref_image_tensor.to(
|
||||
dtype=self.vae.dtype, device=self.vae.device
|
||||
)
|
||||
|
||||
ref_image_latents = self.vae.encode(ref_image_tensor).latent_dist.mean
|
||||
ref_image_latents = ref_image_latents * 0.18215 # (b, 4, h, w)
|
||||
|
||||
gt_image_tensor = self.ref_image_processor.preprocess(
|
||||
gt_image, height=height, width=width) # (bs, c, width, height)
|
||||
gt_image_tensor = gt_image_tensor.to(dtype=self.vae.dtype, device=self.vae.device) # f*c*h*w
|
||||
#do vae for gt image
|
||||
gt_encode = self.vae.encode(gt_image_tensor).latent_dist.mean
|
||||
gt_encode = gt_encode.unsqueeze(2) # b*c*f*h*w
|
||||
gt_encode = gt_encode * 0.18215
|
||||
|
||||
# audio_tensor = torch.zeros((gt_encode.shape[0],50,384))
|
||||
audio_tensor = torch.tensor(audio_tensor).to(dtype=self.denoising_unet.dtype, device=self.denoising_unet.device).unsqueeze(0)
|
||||
uncond_audio_tensor = torch.zeros_like(audio_tensor)
|
||||
audio_tensor = torch.cat([uncond_audio_tensor, audio_tensor], dim=0)
|
||||
audio_tensor = audio_tensor.to(dtype=self.denoising_unet.dtype, device=self.denoising_unet.device)
|
||||
|
||||
# denoising loop
|
||||
num_warmup_steps = len(timesteps) - \
|
||||
num_inference_steps * self.scheduler.order
|
||||
with self.progress_bar(total=num_inference_steps) as progress_bar:
|
||||
for i, t in enumerate(timesteps):
|
||||
# 1. Forward reference image
|
||||
if i == 0:
|
||||
self.reference_unet(
|
||||
ref_image_latents.repeat(
|
||||
(2 if do_classifier_free_guidance else 1), 1, 1, 1
|
||||
),
|
||||
torch.zeros_like(t),
|
||||
encoder_hidden_states=None,
|
||||
return_dict=False,
|
||||
)
|
||||
|
||||
# 2. Update reference unet feature into denosing net
|
||||
reference_control_reader.update(reference_control_writer)
|
||||
|
||||
latents = latents*mask + gt_encode*(1-mask)
|
||||
|
||||
# 3.1 expand the latents if we are doing classifier free guidance
|
||||
latent_model_input = (
|
||||
torch.cat(
|
||||
[latents] * 2) if do_classifier_free_guidance else latents
|
||||
)
|
||||
latent_model_input = self.scheduler.scale_model_input(
|
||||
latent_model_input, t
|
||||
)
|
||||
|
||||
noise_pred = self.denoising_unet(
|
||||
latent_model_input,
|
||||
t,
|
||||
encoder_hidden_states=None,
|
||||
audio_cond_fea=audio_tensor,
|
||||
return_dict=False,
|
||||
)[0]
|
||||
|
||||
# perform guidance
|
||||
if do_classifier_free_guidance:
|
||||
noise_pred_uncond, noise_pred_text = noise_pred.chunk(2)
|
||||
noise_pred = noise_pred_uncond + guidance_scale * (
|
||||
noise_pred_text - noise_pred_uncond
|
||||
)
|
||||
|
||||
# compute the previous noisy sample x_t -> x_t-1
|
||||
latents = self.scheduler.step(
|
||||
noise_pred, t, latents, **extra_step_kwargs, return_dict=False
|
||||
)[0]
|
||||
|
||||
# call the callback, if provided
|
||||
if i == len(timesteps) - 1 or (
|
||||
(i + 1) > num_warmup_steps and (i +
|
||||
1) % self.scheduler.order == 0
|
||||
):
|
||||
progress_bar.update()
|
||||
if callback is not None and i % callback_steps == 0:
|
||||
step_idx = i // getattr(self.scheduler, "order", 1)
|
||||
callback(step_idx, t, latents)
|
||||
reference_control_reader.clear()
|
||||
reference_control_writer.clear()
|
||||
|
||||
latents = latents*mask + gt_encode*(1-mask)
|
||||
|
||||
# Post-processing
|
||||
image = self.decode_latents(latents) # (b, c, 1, h, w)
|
||||
|
||||
# Convert to tensor
|
||||
if output_type == "tensor":
|
||||
image = torch.from_numpy(image)
|
||||
|
||||
if not return_dict:
|
||||
return image
|
||||
|
||||
return StaticPipelineOutput(images=image)
|
||||
@@ -0,0 +1,578 @@
|
||||
import inspect
|
||||
import math
|
||||
from dataclasses import dataclass
|
||||
from typing import Callable, List, Optional, Union
|
||||
|
||||
import numpy as np
|
||||
import torch
|
||||
from diffusers import DiffusionPipeline
|
||||
import torch.nn.functional as F
|
||||
from diffusers.image_processor import VaeImageProcessor
|
||||
from diffusers.schedulers import (
|
||||
DDIMScheduler,
|
||||
DPMSolverMultistepScheduler,
|
||||
EulerAncestralDiscreteScheduler,
|
||||
EulerDiscreteScheduler,
|
||||
LMSDiscreteScheduler,
|
||||
PNDMScheduler,
|
||||
)
|
||||
from diffusers.utils import BaseOutput, deprecate, is_accelerate_available, logging
|
||||
from diffusers.utils.torch_utils import randn_tensor
|
||||
from einops import rearrange
|
||||
from tqdm import tqdm
|
||||
from transformers import CLIPImageProcessor
|
||||
|
||||
from ..models.mutual_self_attention import ReferenceAttentionControl
|
||||
from .context import get_context_scheduler
|
||||
from .utils import get_tensor_interpolation_method
|
||||
|
||||
@dataclass
|
||||
class Audio2VideoPipelineOutput(BaseOutput):
|
||||
videos: Union[torch.Tensor, np.ndarray]
|
||||
|
||||
|
||||
class Audio2VideoPipeline(DiffusionPipeline):
|
||||
_optional_components = []
|
||||
|
||||
def __init__(
|
||||
self,
|
||||
vae,
|
||||
reference_unet,
|
||||
denoising_unet,
|
||||
audio_guider,
|
||||
face_locator,
|
||||
scheduler: Union[
|
||||
DDIMScheduler,
|
||||
PNDMScheduler,
|
||||
LMSDiscreteScheduler,
|
||||
EulerDiscreteScheduler,
|
||||
EulerAncestralDiscreteScheduler,
|
||||
DPMSolverMultistepScheduler,
|
||||
],
|
||||
image_proj_model=None,
|
||||
tokenizer=None,
|
||||
text_encoder=None,
|
||||
):
|
||||
super().__init__()
|
||||
|
||||
self.register_modules(
|
||||
vae=vae,
|
||||
reference_unet=reference_unet,
|
||||
denoising_unet=denoising_unet,
|
||||
audio_guider=audio_guider,
|
||||
face_locator=face_locator,
|
||||
scheduler=scheduler,
|
||||
image_proj_model=image_proj_model,
|
||||
tokenizer=tokenizer,
|
||||
text_encoder=text_encoder,
|
||||
)
|
||||
self.vae_scale_factor = 2 ** (len(self.vae.config.block_out_channels) - 1)
|
||||
self.clip_image_processor = CLIPImageProcessor()
|
||||
self.ref_image_processor = VaeImageProcessor(
|
||||
vae_scale_factor=self.vae_scale_factor, do_convert_rgb=True
|
||||
)
|
||||
|
||||
def enable_vae_slicing(self):
|
||||
self.vae.enable_slicing()
|
||||
|
||||
def disable_vae_slicing(self):
|
||||
self.vae.disable_slicing()
|
||||
|
||||
def enable_sequential_cpu_offload(self, gpu_id=0):
|
||||
if is_accelerate_available():
|
||||
from accelerate import cpu_offload
|
||||
else:
|
||||
raise ImportError("Please install accelerate via `pip install accelerate`")
|
||||
|
||||
device = torch.device(f"cuda:{gpu_id}")
|
||||
|
||||
for cpu_offloaded_model in [self.unet, self.text_encoder, self.vae]:
|
||||
if cpu_offloaded_model is not None:
|
||||
cpu_offload(cpu_offloaded_model, device)
|
||||
|
||||
@property
|
||||
def _execution_device(self):
|
||||
if self.device != torch.device("meta") or not hasattr(self.unet, "_hf_hook"):
|
||||
return self.device
|
||||
for module in self.unet.modules():
|
||||
if (
|
||||
hasattr(module, "_hf_hook")
|
||||
and hasattr(module._hf_hook, "execution_device")
|
||||
and module._hf_hook.execution_device is not None
|
||||
):
|
||||
return torch.device(module._hf_hook.execution_device)
|
||||
return self.device
|
||||
|
||||
def decode_latents(self, latents):
|
||||
video_length = latents.shape[2]
|
||||
latents = 1 / 0.18215 * latents
|
||||
latents = rearrange(latents, "b c f h w -> (b f) c h w")
|
||||
video = []
|
||||
for frame_idx in tqdm(range(latents.shape[0])):
|
||||
video.append(self.vae.decode(latents[frame_idx : frame_idx + 1]).sample)
|
||||
video = torch.cat(video)
|
||||
video = rearrange(video, "(b f) c h w -> b c f h w", f=video_length)
|
||||
video = (video / 2 + 0.5).clamp(0, 1)
|
||||
# we always cast to float32 as this does not cause significant overhead and is compatible with bfloa16
|
||||
video = video.cpu().float().numpy()
|
||||
return video
|
||||
|
||||
def prepare_extra_step_kwargs(self, generator, eta):
|
||||
# prepare extra kwargs for the scheduler step, since not all schedulers have the same signature
|
||||
# eta (η) is only used with the DDIMScheduler, it will be ignored for other schedulers.
|
||||
# eta corresponds to η in DDIM paper: https://arxiv.org/abs/2010.02502
|
||||
# and should be between [0, 1]
|
||||
|
||||
accepts_eta = "eta" in set(
|
||||
inspect.signature(self.scheduler.step).parameters.keys()
|
||||
)
|
||||
extra_step_kwargs = {}
|
||||
if accepts_eta:
|
||||
extra_step_kwargs["eta"] = eta
|
||||
|
||||
# check if the scheduler accepts generator
|
||||
accepts_generator = "generator" in set(
|
||||
inspect.signature(self.scheduler.step).parameters.keys()
|
||||
)
|
||||
if accepts_generator:
|
||||
extra_step_kwargs["generator"] = generator
|
||||
return extra_step_kwargs
|
||||
|
||||
def prepare_latents(
|
||||
self,
|
||||
batch_size,
|
||||
num_channels_latents,
|
||||
width,
|
||||
height,
|
||||
video_length,
|
||||
dtype,
|
||||
device,
|
||||
generator
|
||||
):
|
||||
shape = (
|
||||
batch_size,
|
||||
num_channels_latents,
|
||||
# context_frame_length,
|
||||
video_length,
|
||||
height // self.vae_scale_factor,
|
||||
width // self.vae_scale_factor,
|
||||
)
|
||||
|
||||
if isinstance(generator, list) and len(generator) != batch_size:
|
||||
raise ValueError(
|
||||
f"You have passed a list of generators of length {len(generator)}, but requested an effective batch"
|
||||
f" size of {batch_size}. Make sure the batch size matches the length of the generators."
|
||||
)
|
||||
|
||||
latents_seg = randn_tensor(
|
||||
shape, generator=generator, device=device, dtype=dtype
|
||||
)
|
||||
latents = latents_seg
|
||||
latents = torch.clamp(latents, -1.5, 1.5)
|
||||
latents = latents * self.scheduler.init_noise_sigma
|
||||
return latents
|
||||
|
||||
def _encode_prompt(
|
||||
self,
|
||||
prompt,
|
||||
device,
|
||||
num_videos_per_prompt,
|
||||
do_classifier_free_guidance,
|
||||
negative_prompt,
|
||||
):
|
||||
batch_size = len(prompt) if isinstance(prompt, list) else 1
|
||||
|
||||
text_inputs = self.tokenizer(
|
||||
prompt,
|
||||
padding="max_length",
|
||||
max_length=self.tokenizer.model_max_length,
|
||||
truncation=True,
|
||||
return_tensors="pt",
|
||||
)
|
||||
text_input_ids = text_inputs.input_ids
|
||||
untruncated_ids = self.tokenizer(
|
||||
prompt, padding="longest", return_tensors="pt"
|
||||
).input_ids
|
||||
|
||||
if untruncated_ids.shape[-1] >= text_input_ids.shape[-1] and not torch.equal(
|
||||
text_input_ids, untruncated_ids
|
||||
):
|
||||
removed_text = self.tokenizer.batch_decode(
|
||||
untruncated_ids[:, self.tokenizer.model_max_length - 1 : -1]
|
||||
)
|
||||
|
||||
if (
|
||||
hasattr(self.text_encoder.config, "use_attention_mask")
|
||||
and self.text_encoder.config.use_attention_mask
|
||||
):
|
||||
attention_mask = text_inputs.attention_mask.to(device)
|
||||
else:
|
||||
attention_mask = None
|
||||
|
||||
text_embeddings = self.text_encoder(
|
||||
text_input_ids.to(device),
|
||||
attention_mask=attention_mask,
|
||||
)
|
||||
text_embeddings = text_embeddings[0]
|
||||
|
||||
# duplicate text embeddings for each generation per prompt, using mps friendly method
|
||||
bs_embed, seq_len, _ = text_embeddings.shape
|
||||
text_embeddings = text_embeddings.repeat(1, num_videos_per_prompt, 1)
|
||||
text_embeddings = text_embeddings.view(
|
||||
bs_embed * num_videos_per_prompt, seq_len, -1
|
||||
)
|
||||
|
||||
# get unconditional embeddings for classifier free guidance
|
||||
if do_classifier_free_guidance:
|
||||
uncond_tokens: List[str]
|
||||
if negative_prompt is None:
|
||||
uncond_tokens = [""] * batch_size
|
||||
elif type(prompt) is not type(negative_prompt):
|
||||
raise TypeError(
|
||||
f"`negative_prompt` should be the same type to `prompt`, but got {type(negative_prompt)} !="
|
||||
f" {type(prompt)}."
|
||||
)
|
||||
elif isinstance(negative_prompt, str):
|
||||
uncond_tokens = [negative_prompt]
|
||||
elif batch_size != len(negative_prompt):
|
||||
raise ValueError(
|
||||
f"`negative_prompt`: {negative_prompt} has batch size {len(negative_prompt)}, but `prompt`:"
|
||||
f" {prompt} has batch size {batch_size}. Please make sure that passed `negative_prompt` matches"
|
||||
" the batch size of `prompt`."
|
||||
)
|
||||
else:
|
||||
uncond_tokens = negative_prompt
|
||||
|
||||
max_length = text_input_ids.shape[-1]
|
||||
uncond_input = self.tokenizer(
|
||||
uncond_tokens,
|
||||
padding="max_length",
|
||||
max_length=max_length,
|
||||
truncation=True,
|
||||
return_tensors="pt",
|
||||
)
|
||||
|
||||
if (
|
||||
hasattr(self.text_encoder.config, "use_attention_mask")
|
||||
and self.text_encoder.config.use_attention_mask
|
||||
):
|
||||
attention_mask = uncond_input.attention_mask.to(device)
|
||||
else:
|
||||
attention_mask = None
|
||||
|
||||
uncond_embeddings = self.text_encoder(
|
||||
uncond_input.input_ids.to(device),
|
||||
attention_mask=attention_mask,
|
||||
)
|
||||
uncond_embeddings = uncond_embeddings[0]
|
||||
|
||||
# duplicate unconditional embeddings for each generation per prompt, using mps friendly method
|
||||
seq_len = uncond_embeddings.shape[1]
|
||||
uncond_embeddings = uncond_embeddings.repeat(1, num_videos_per_prompt, 1)
|
||||
uncond_embeddings = uncond_embeddings.view(
|
||||
batch_size * num_videos_per_prompt, seq_len, -1
|
||||
)
|
||||
|
||||
# For classifier free guidance, we need to do two forward passes.
|
||||
# Here we concatenate the unconditional and text embeddings into a single batch
|
||||
# to avoid doing two forward passes
|
||||
text_embeddings = torch.cat([uncond_embeddings, text_embeddings])
|
||||
|
||||
return text_embeddings
|
||||
|
||||
def interpolate_latents(
|
||||
self, latents: torch.Tensor, interpolation_factor: int, device
|
||||
):
|
||||
if interpolation_factor < 2:
|
||||
return latents
|
||||
|
||||
new_latents = torch.zeros(
|
||||
(
|
||||
latents.shape[0],
|
||||
latents.shape[1],
|
||||
((latents.shape[2] - 1) * interpolation_factor) + 1,
|
||||
latents.shape[3],
|
||||
latents.shape[4],
|
||||
),
|
||||
device=latents.device,
|
||||
dtype=latents.dtype,
|
||||
)
|
||||
|
||||
org_video_length = latents.shape[2]
|
||||
rate = [i / interpolation_factor for i in range(interpolation_factor)][1:]
|
||||
|
||||
new_index = 0
|
||||
|
||||
v0 = None
|
||||
v1 = None
|
||||
|
||||
for i0, i1 in zip(range(org_video_length), range(org_video_length)[1:]):
|
||||
v0 = latents[:, :, i0, :, :]
|
||||
v1 = latents[:, :, i1, :, :]
|
||||
|
||||
new_latents[:, :, new_index, :, :] = v0
|
||||
new_index += 1
|
||||
|
||||
for f in rate:
|
||||
v = get_tensor_interpolation_method()(
|
||||
v0.to(device=device), v1.to(device=device), f
|
||||
)
|
||||
new_latents[:, :, new_index, :, :] = v.to(latents.device)
|
||||
new_index += 1
|
||||
|
||||
new_latents[:, :, new_index, :, :] = v1
|
||||
new_index += 1
|
||||
|
||||
return new_latents
|
||||
|
||||
@torch.no_grad()
|
||||
def __call__(
|
||||
self,
|
||||
ref_image,
|
||||
audio_path,
|
||||
face_mask_tensor,
|
||||
width,
|
||||
height,
|
||||
video_length,
|
||||
num_inference_steps,
|
||||
guidance_scale,
|
||||
num_images_per_prompt=1,
|
||||
eta: float = 0.0,
|
||||
generator: Optional[Union[torch.Generator, List[torch.Generator]]] = None,
|
||||
output_type: Optional[str] = "tensor",
|
||||
return_dict: bool = True,
|
||||
callback: Optional[Callable[[int, int, torch.FloatTensor], None]] = None,
|
||||
callback_steps: Optional[int] = 1,
|
||||
context_schedule="uniform",
|
||||
context_frames=12,
|
||||
context_stride=1,
|
||||
context_overlap=0,
|
||||
context_batch_size=1,
|
||||
interpolation_factor=1,
|
||||
audio_sample_rate=16000,
|
||||
fps=25,
|
||||
audio_margin=2,
|
||||
**kwargs,
|
||||
):
|
||||
# Default height and width to unet
|
||||
height = height or self.unet.config.sample_size * self.vae_scale_factor
|
||||
width = width or self.unet.config.sample_size * self.vae_scale_factor
|
||||
|
||||
device = self._execution_device
|
||||
|
||||
do_classifier_free_guidance = guidance_scale > 1.0
|
||||
|
||||
# Prepare timesteps
|
||||
self.scheduler.set_timesteps(num_inference_steps, device=device)
|
||||
timesteps = self.scheduler.timesteps
|
||||
|
||||
batch_size = 1
|
||||
|
||||
reference_control_writer = ReferenceAttentionControl(
|
||||
self.reference_unet,
|
||||
do_classifier_free_guidance=do_classifier_free_guidance,
|
||||
mode="write",
|
||||
batch_size=batch_size,
|
||||
fusion_blocks="full",
|
||||
)
|
||||
reference_control_reader = ReferenceAttentionControl(
|
||||
self.denoising_unet,
|
||||
do_classifier_free_guidance=do_classifier_free_guidance,
|
||||
mode="read",
|
||||
batch_size=batch_size,
|
||||
fusion_blocks="full",
|
||||
)
|
||||
|
||||
whisper_feature = self.audio_guider.audio2feat(audio_path)
|
||||
|
||||
whisper_chunks = self.audio_guider.feature2chunks(feature_array=whisper_feature, fps=fps)
|
||||
|
||||
print("whisper_chunks:", whisper_chunks.shape)
|
||||
audio_frame_num = whisper_chunks.shape[0]
|
||||
audio_fea_final = torch.Tensor(whisper_chunks).to(dtype=self.vae.dtype, device=self.vae.device)
|
||||
audio_fea_final = audio_fea_final.unsqueeze(0)
|
||||
print("audio_fea_final:", audio_fea_final.shape)
|
||||
video_length = min(video_length, audio_frame_num)
|
||||
if video_length < audio_frame_num:
|
||||
audio_fea_final = audio_fea_final[:, :video_length, :, :]
|
||||
|
||||
num_channels_latents = self.denoising_unet.in_channels
|
||||
latents = self.prepare_latents(
|
||||
batch_size * num_images_per_prompt,
|
||||
num_channels_latents,
|
||||
width,
|
||||
height,
|
||||
video_length,
|
||||
audio_fea_final.dtype,
|
||||
device,
|
||||
generator
|
||||
)
|
||||
# print(video_length, latents.shape)
|
||||
c_face_locator_tensor = self.face_locator(face_mask_tensor)
|
||||
uc_face_locator_tensor = torch.zeros_like(c_face_locator_tensor)
|
||||
face_locator_tensor = torch.cat([uc_face_locator_tensor, c_face_locator_tensor], dim=0)
|
||||
# Prepare extra step kwargs.
|
||||
extra_step_kwargs = self.prepare_extra_step_kwargs(generator, eta)
|
||||
|
||||
# Prepare ref image latents
|
||||
ref_image_tensor = self.ref_image_processor.preprocess(
|
||||
ref_image, height=height, width=width
|
||||
)
|
||||
ref_image_tensor = ref_image_tensor.to(
|
||||
dtype=self.vae.dtype, device=self.vae.device
|
||||
)
|
||||
ref_image_latents = self.vae.encode(ref_image_tensor).latent_dist.mean
|
||||
ref_image_latents = ref_image_latents * 0.18215 # (b , 4, h, w)
|
||||
|
||||
context_scheduler = get_context_scheduler(context_schedule)
|
||||
|
||||
# denoising loop
|
||||
num_warmup_steps = len(timesteps) - num_inference_steps * self.scheduler.order
|
||||
context_queue = list(
|
||||
context_scheduler(
|
||||
0,
|
||||
num_inference_steps,
|
||||
latents.shape[2],
|
||||
context_frames,
|
||||
context_stride,
|
||||
context_overlap,
|
||||
)
|
||||
)
|
||||
print("ref_image_latents shape:", ref_image_latents.shape)
|
||||
print("face_mask_tensor shape:", face_mask_tensor.shape)
|
||||
print("face_locator_tensor shape:", face_locator_tensor.shape)
|
||||
with self.progress_bar(total=num_inference_steps) as progress_bar:
|
||||
for t_i, t in enumerate(timesteps):
|
||||
noise_pred = torch.zeros(
|
||||
(
|
||||
latents.shape[0] * (2 if do_classifier_free_guidance else 1),
|
||||
*latents.shape[1:],
|
||||
),
|
||||
device=latents.device,
|
||||
dtype=latents.dtype,
|
||||
)
|
||||
counter = torch.zeros(
|
||||
(1, 1, latents.shape[2], 1, 1),
|
||||
device=latents.device,
|
||||
dtype=latents.dtype,
|
||||
)
|
||||
|
||||
# 1. Forward reference image
|
||||
if t_i == 0:
|
||||
self.reference_unet(
|
||||
ref_image_latents,
|
||||
torch.zeros_like(t),
|
||||
encoder_hidden_states=None,
|
||||
return_dict=False,
|
||||
)
|
||||
reference_control_reader.update(reference_control_writer, do_classifier_free_guidance=do_classifier_free_guidance)
|
||||
|
||||
|
||||
num_context_batches = math.ceil(len(context_queue) / context_batch_size)
|
||||
|
||||
global_context = []
|
||||
for j in range(num_context_batches):
|
||||
global_context.append(
|
||||
context_queue[
|
||||
j * context_batch_size : (j + 1) * context_batch_size
|
||||
]
|
||||
)
|
||||
|
||||
for context in global_context:
|
||||
new_context = [[0 for _ in range(len(context[c_j]))] for c_j in range(len(context))]
|
||||
for c_j in range(len(context)):
|
||||
for c_i in range(len(context[c_j])):
|
||||
new_context[c_j][c_i] = (context[c_j][c_i] + t_i * 2) % video_length
|
||||
|
||||
latent_model_input = (
|
||||
torch.cat([latents[:, :, c] for c in new_context])
|
||||
.to(device)
|
||||
.repeat(2 if do_classifier_free_guidance else 1, 1, 1, 1, 1)
|
||||
)
|
||||
c_audio_latents = torch.cat([audio_fea_final[:, c] for c in new_context]).to(device)
|
||||
audio_latents = torch.cat([torch.zeros_like(c_audio_latents), c_audio_latents], 0)
|
||||
|
||||
latent_model_input = self.scheduler.scale_model_input(
|
||||
latent_model_input, t
|
||||
)
|
||||
print(c_audio_latents.shape)
|
||||
pred = self.denoising_unet(
|
||||
latent_model_input,
|
||||
t,
|
||||
encoder_hidden_states=None,
|
||||
audio_cond_fea=audio_latents if do_classifier_free_guidance else c_audio_latents,
|
||||
face_musk_fea=face_locator_tensor if do_classifier_free_guidance else c_face_locator_tensor,
|
||||
return_dict=False,
|
||||
)[0]
|
||||
|
||||
for j, c in enumerate(new_context):
|
||||
noise_pred[:, :, c] = noise_pred[:, :, c] + pred
|
||||
counter[:, :, c] = counter[:, :, c] + 1
|
||||
|
||||
# perform guidance
|
||||
if do_classifier_free_guidance:
|
||||
noise_pred_uncond, noise_pred_text = (noise_pred / counter).chunk(2)
|
||||
noise_pred = noise_pred_uncond + guidance_scale * (
|
||||
noise_pred_text - noise_pred_uncond
|
||||
)
|
||||
|
||||
latents = self.scheduler.step(
|
||||
noise_pred, t, latents, **extra_step_kwargs
|
||||
).prev_sample
|
||||
|
||||
if t_i == len(timesteps) - 1 or (
|
||||
(t_i + 1) > num_warmup_steps and (t_i + 1) % self.scheduler.order == 0
|
||||
):
|
||||
progress_bar.update()
|
||||
|
||||
reference_control_reader.clear()
|
||||
reference_control_writer.clear()
|
||||
|
||||
if interpolation_factor > 0:
|
||||
latents = self.interpolate_latents(latents, interpolation_factor, device)
|
||||
# Post-processing
|
||||
images = self.decode_latents(latents) # (b, c, f, h, w)
|
||||
|
||||
# Convert to tensor
|
||||
if output_type == "tensor":
|
||||
images = torch.from_numpy(images)
|
||||
|
||||
if not return_dict:
|
||||
return images
|
||||
|
||||
return Audio2VideoPipelineOutput(videos=images)
|
||||
|
||||
def smooth_f_axis(self, tensor, smoothing_coef=0.2):
|
||||
"""
|
||||
对5维tensor的F轴进行平滑，首尾帧保持原始值，其他帧应用平滑操作。
|
||||
|
||||
参数:
|
||||
tensor -- 输入的5D Tensor，形状为(B, C, F, W, H)
|
||||
smoothing_coef -- 平滑前后帧的系数，默认为0.3
|
||||
|
||||
返回:
|
||||
smoothed_tensor -- 平滑后的5D Tensor
|
||||
"""
|
||||
# 生成平滑核
|
||||
weight_kernel = torch.tensor([smoothing_coef, 1 - 2 * smoothing_coef, smoothing_coef],
|
||||
device=tensor.device, dtype=tensor.dtype).view(1, 1, -1)
|
||||
|
||||
b, c, f, w, h = tensor.size()
|
||||
|
||||
# 对于第三轴（F轴）仅平滑内部帧 (1:-1)
|
||||
# 首先，移动要平滑的轴到最后一个维度以便应用conv1d
|
||||
tensor_moved = rearrange(tensor, "b c f h w -> b (c h w) f")
|
||||
|
||||
# 对除了两端帧外的所有帧应用conv1d
|
||||
internal_frames = F.conv1d(tensor_moved, weight_kernel)
|
||||
print("tensor:", tensor.shape)
|
||||
|
||||
# 重新整理输出形状为 (B, F, C, W, H)
|
||||
internal_frames = rearrange(internal_frames, "b (c h w) f -> b c f h w", c=c, h=h, w=w)
|
||||
print("internal_frames:", internal_frames.shape)
|
||||
# 将第一帧和最后一帧保持不变，合并结果
|
||||
# 首帧 tensor[:, :, 0:1, :, :], 中间帧 internal_frames[:, :, 1:-1, :, :], 最后帧 tensor[:, :, -1:, :, :]
|
||||
smoothed_tensor = torch.cat(
|
||||
[tensor[:, :, 0:1, :, :], internal_frames, tensor[:, :, -1:, :, :]], dim=2)
|
||||
|
||||
return smoothed_tensor
|
||||
@@ -0,0 +1,29 @@
|
||||
import torch
|
||||
|
||||
tensor_interpolation = None
|
||||
|
||||
|
||||
def get_tensor_interpolation_method():
|
||||
return tensor_interpolation
|
||||
|
||||
|
||||
def set_tensor_interpolation_method(is_slerp):
|
||||
global tensor_interpolation
|
||||
tensor_interpolation = slerp if is_slerp else linear
|
||||
|
||||
|
||||
def linear(v1, v2, t):
|
||||
return (1.0 - t) * v1 + t * v2
|
||||
|
||||
|
||||
def slerp(
|
||||
v0: torch.Tensor, v1: torch.Tensor, t: float, DOT_THRESHOLD: float = 0.9995
|
||||
) -> torch.Tensor:
|
||||
u0 = v0 / v0.norm()
|
||||
u1 = v1 / v1.norm()
|
||||
dot = (u0 * u1).sum()
|
||||
if dot.abs() > DOT_THRESHOLD:
|
||||
# logger.info(f'warning: v0 and v1 close to parallel, using linear interpolation instead.')
|
||||
return (1.0 - t) * v0 + t * v1
|
||||
omega = dot.acos()
|
||||
return (((1.0 - t) * omega).sin() * v0 + (t * omega).sin() * v1) / omega.sin()
|
||||
@@ -0,0 +1,190 @@
|
||||
import cv2
|
||||
import mediapipe as mp
|
||||
import numpy as np
|
||||
from mediapipe.framework.formats import landmark_pb2
|
||||
|
||||
class FaceMeshVisualizer:
|
||||
def __init__(self,
|
||||
forehead_edge=False,
|
||||
upface_only=False,
|
||||
draw_eye=True,
|
||||
draw_head=False,
|
||||
draw_iris=True,
|
||||
draw_eyebrow=True,
|
||||
draw_mouse=True,
|
||||
draw_nose=True,
|
||||
draw_pupil=True
|
||||
):
|
||||
self.mp_drawing = mp.solutions.drawing_utils
|
||||
mp_face_mesh = mp.solutions.face_mesh
|
||||
self.mp_face_mesh = mp_face_mesh
|
||||
self.forehead_edge = forehead_edge
|
||||
|
||||
DrawingSpec = mp.solutions.drawing_styles.DrawingSpec
|
||||
f_thick = 2
|
||||
f_rad = 1
|
||||
right_iris_draw = DrawingSpec(color=(10, 200, 250), thickness=f_thick, circle_radius=f_rad)
|
||||
right_eye_draw = DrawingSpec(color=(10, 200, 180), thickness=f_thick, circle_radius=f_rad)
|
||||
right_eyebrow_draw = DrawingSpec(color=(10, 220, 180), thickness=f_thick, circle_radius=f_rad)
|
||||
left_iris_draw = DrawingSpec(color=(250, 200, 10), thickness=f_thick, circle_radius=f_rad)
|
||||
left_eye_draw = DrawingSpec(color=(180, 200, 10), thickness=f_thick, circle_radius=f_rad)
|
||||
left_eyebrow_draw = DrawingSpec(color=(180, 220, 10), thickness=f_thick, circle_radius=f_rad)
|
||||
head_draw = DrawingSpec(color=(10, 200, 10), thickness=f_thick, circle_radius=f_rad)
|
||||
nose_draw = DrawingSpec(color=(200, 200, 200), thickness=f_thick, circle_radius=f_rad)
|
||||
|
||||
mouth_draw_obl = DrawingSpec(color=(10, 180, 20), thickness=f_thick, circle_radius=f_rad)
|
||||
mouth_draw_obr = DrawingSpec(color=(20, 10, 180), thickness=f_thick, circle_radius=f_rad)
|
||||
|
||||
mouth_draw_ibl = DrawingSpec(color=(100, 100, 30), thickness=f_thick, circle_radius=f_rad)
|
||||
mouth_draw_ibr = DrawingSpec(color=(100, 150, 50), thickness=f_thick, circle_radius=f_rad)
|
||||
|
||||
mouth_draw_otl = DrawingSpec(color=(20, 80, 100), thickness=f_thick, circle_radius=f_rad)
|
||||
mouth_draw_otr = DrawingSpec(color=(80, 100, 20), thickness=f_thick, circle_radius=f_rad)
|
||||
|
||||
mouth_draw_itl = DrawingSpec(color=(120, 100, 200), thickness=f_thick, circle_radius=f_rad)
|
||||
mouth_draw_itr = DrawingSpec(color=(150 ,120, 100), thickness=f_thick, circle_radius=f_rad)
|
||||
|
||||
FACEMESH_LIPS_OUTER_BOTTOM_LEFT = [(61,146),(146,91),(91,181),(181,84),(84,17)]
|
||||
FACEMESH_LIPS_OUTER_BOTTOM_RIGHT = [(17,314),(314,405),(405,321),(321,375),(375,291)]
|
||||
|
||||
FACEMESH_LIPS_INNER_BOTTOM_LEFT = [(78,95),(95,88),(88,178),(178,87),(87,14)]
|
||||
FACEMESH_LIPS_INNER_BOTTOM_RIGHT = [(14,317),(317,402),(402,318),(318,324),(324,308)]
|
||||
|
||||
FACEMESH_LIPS_OUTER_TOP_LEFT = [(61,185),(185,40),(40,39),(39,37),(37,0)]
|
||||
FACEMESH_LIPS_OUTER_TOP_RIGHT = [(0,267),(267,269),(269,270),(270,409),(409,291)]
|
||||
|
||||
FACEMESH_LIPS_INNER_TOP_LEFT = [(78,191),(191,80),(80,81),(81,82),(82,13)]
|
||||
FACEMESH_LIPS_INNER_TOP_RIGHT = [(13,312),(312,311),(311,310),(310,415),(415,308)]
|
||||
|
||||
FACEMESH_CUSTOM_FACE_OVAL = [(176, 149), (150, 136), (356, 454), (58, 132), (152, 148), (361, 288), (251, 389), (132, 93), (389, 356), (400, 377), (136, 172), (377, 152), (323, 361), (172, 58), (454, 323), (365, 379), (379, 378), (148, 176), (93, 234), (397, 365), (149, 150), (288, 397), (234, 127), (378, 400), (127, 162), (162, 21)]
|
||||
|
||||
# mp_face_mesh.FACEMESH_CONTOURS has all the items we care about.
|
||||
face_connection_spec = {}
|
||||
|
||||
#from IPython import embed
|
||||
#embed()
|
||||
if self.forehead_edge:
|
||||
for edge in mp_face_mesh.FACEMESH_FACE_OVAL:
|
||||
face_connection_spec[edge] = head_draw
|
||||
else:
|
||||
if draw_head:
|
||||
FACEMESH_CUSTOM_FACE_OVAL_sorted = sorted(FACEMESH_CUSTOM_FACE_OVAL)
|
||||
if upface_only:
|
||||
for edge in [FACEMESH_CUSTOM_FACE_OVAL_sorted[edge_idx] for edge_idx in [1,2,9,12,13,16,22,25]]:
|
||||
face_connection_spec[edge] = head_draw
|
||||
else:
|
||||
for edge in FACEMESH_CUSTOM_FACE_OVAL_sorted:
|
||||
face_connection_spec[edge] = head_draw
|
||||
|
||||
if draw_eye:
|
||||
for edge in mp_face_mesh.FACEMESH_LEFT_EYE:
|
||||
face_connection_spec[edge] = left_eye_draw
|
||||
for edge in mp_face_mesh.FACEMESH_RIGHT_EYE:
|
||||
face_connection_spec[edge] = right_eye_draw
|
||||
|
||||
if draw_eyebrow:
|
||||
for edge in mp_face_mesh.FACEMESH_LEFT_EYEBROW:
|
||||
face_connection_spec[edge] = left_eyebrow_draw
|
||||
|
||||
for edge in mp_face_mesh.FACEMESH_RIGHT_EYEBROW:
|
||||
face_connection_spec[edge] = right_eyebrow_draw
|
||||
|
||||
if draw_iris:
|
||||
for edge in mp_face_mesh.FACEMESH_LEFT_IRIS:
|
||||
face_connection_spec[edge] = left_iris_draw
|
||||
for edge in mp_face_mesh.FACEMESH_RIGHT_IRIS:
|
||||
face_connection_spec[edge] = right_iris_draw
|
||||
|
||||
#for edge in mp_face_mesh.FACEMESH_RIGHT_EYEBROW:
|
||||
# face_connection_spec[edge] = right_eyebrow_draw
|
||||
# for edge in mp_face_mesh.FACEMESH_RIGHT_IRIS:
|
||||
# face_connection_spec[edge] = right_iris_draw
|
||||
|
||||
# for edge in mp_face_mesh.FACEMESH_LIPS:
|
||||
# face_connection_spec[edge] = mouth_draw
|
||||
|
||||
if draw_mouse:
|
||||
for edge in FACEMESH_LIPS_OUTER_BOTTOM_LEFT:
|
||||
face_connection_spec[edge] = mouth_draw_obl
|
||||
for edge in FACEMESH_LIPS_OUTER_BOTTOM_RIGHT:
|
||||
face_connection_spec[edge] = mouth_draw_obr
|
||||
for edge in FACEMESH_LIPS_INNER_BOTTOM_LEFT:
|
||||
face_connection_spec[edge] = mouth_draw_ibl
|
||||
for edge in FACEMESH_LIPS_INNER_BOTTOM_RIGHT:
|
||||
face_connection_spec[edge] = mouth_draw_ibr
|
||||
for edge in FACEMESH_LIPS_OUTER_TOP_LEFT:
|
||||
face_connection_spec[edge] = mouth_draw_otl
|
||||
for edge in FACEMESH_LIPS_OUTER_TOP_RIGHT:
|
||||
face_connection_spec[edge] = mouth_draw_otr
|
||||
for edge in FACEMESH_LIPS_INNER_TOP_LEFT:
|
||||
face_connection_spec[edge] = mouth_draw_itl
|
||||
for edge in FACEMESH_LIPS_INNER_TOP_RIGHT:
|
||||
face_connection_spec[edge] = mouth_draw_itr
|
||||
|
||||
self.face_connection_spec = face_connection_spec
|
||||
|
||||
self.pupil_landmark_spec = {468: right_iris_draw, 473: left_iris_draw}
|
||||
self.nose_landmark_spec = {4: nose_draw}
|
||||
|
||||
self.draw_pupil = draw_pupil
|
||||
self.draw_nose = draw_nose
|
||||
|
||||
def draw_points(self, image, landmark_list, drawing_spec, halfwidth: int = 2):
|
||||
"""We have a custom function to draw the pupils because the mp.draw_landmarks method requires a parameter for all
|
||||
landmarks. Until our PR is merged into mediapipe, we need this separate method."""
|
||||
if len(image.shape) != 3:
|
||||
raise ValueError("Input image must be H,W,C.")
|
||||
image_rows, image_cols, image_channels = image.shape
|
||||
if image_channels != 3: # BGR channels
|
||||
raise ValueError('Input image must contain three channel bgr data.')
|
||||
for idx, landmark in enumerate(landmark_list.landmark):
|
||||
if idx not in drawing_spec:
|
||||
continue
|
||||
|
||||
if (
|
||||
(landmark.HasField('visibility') and landmark.visibility < 0.9) or
|
||||
(landmark.HasField('presence') and landmark.presence < 0.5)
|
||||
):
|
||||
continue
|
||||
if landmark.x >= 1.0 or landmark.x < 0 or landmark.y >= 1.0 or landmark.y < 0:
|
||||
continue
|
||||
|
||||
image_x = int(image_cols * landmark.x)
|
||||
image_y = int(image_rows * landmark.y)
|
||||
|
||||
draw_color = drawing_spec[idx].color
|
||||
image[image_y - halfwidth : image_y + halfwidth, image_x - halfwidth : image_x + halfwidth, :] = draw_color
|
||||
|
||||
|
||||
def draw_landmarks(self, image_size, keypoints, normed=False):
|
||||
ini_size = [512, 512]
|
||||
image = np.zeros([ini_size[1], ini_size[0], 3], dtype=np.uint8)
|
||||
if keypoints is not None:
|
||||
new_landmarks = landmark_pb2.NormalizedLandmarkList()
|
||||
for i in range(keypoints.shape[0]):
|
||||
landmark = new_landmarks.landmark.add()
|
||||
if normed:
|
||||
landmark.x = keypoints[i, 0]
|
||||
landmark.y = keypoints[i, 1]
|
||||
else:
|
||||
landmark.x = keypoints[i, 0] / image_size[0]
|
||||
landmark.y = keypoints[i, 1] / image_size[1]
|
||||
landmark.z = 1.0
|
||||
|
||||
self.mp_drawing.draw_landmarks(
|
||||
image=image,
|
||||
landmark_list=new_landmarks,
|
||||
connections=self.face_connection_spec.keys(),
|
||||
landmark_drawing_spec=None,
|
||||
connection_drawing_spec=self.face_connection_spec
|
||||
)
|
||||
|
||||
if self.draw_pupil:
|
||||
self.draw_points(image, new_landmarks, self.pupil_landmark_spec, 3)
|
||||
|
||||
if self.draw_nose:
|
||||
self.draw_points(image, new_landmarks, self.nose_landmark_spec, 3)
|
||||
|
||||
image = cv2.resize(image, (image_size[0], image_size[1]))
|
||||
|
||||
return image
|
||||
File diff suppressed because it is too large
Load Diff
@@ -0,0 +1,84 @@
|
||||
from PIL import Image
|
||||
import cv2
|
||||
import numpy as np
|
||||
from imageio_ffmpeg import get_ffmpeg_exe
|
||||
import pathlib
|
||||
import os
|
||||
from IPython import embed
|
||||
|
||||
def pil_to_cv2(pil):
|
||||
return cv2.cvtColor(np.array(pil).astype(np.uint8), cv2.COLOR_RGB2BGR)
|
||||
|
||||
def cv2_to_pil(cv2_img):
|
||||
return Image.fromarray(cv2.cvtColor(cv2_img, cv2.COLOR_BGR2RGB).astype(np.uint8))
|
||||
|
||||
def center_crop_cv2(cv2_pic):
|
||||
h, w = cv2_pic.shape[0], cv2_pic.shape[1]
|
||||
if h > w:
|
||||
return cv2_pic[(h - w) // 2 : (h - w) // 2 + w, :]
|
||||
else:
|
||||
return cv2_pic[:, (w - h) // 2 : (w - h) // 2 + h]
|
||||
|
||||
def pils_from_video(video_path):
|
||||
cap = cv2.VideoCapture(video_path)
|
||||
pils = []
|
||||
while True:
|
||||
ret, frame = cap.read()
|
||||
if not ret:
|
||||
break
|
||||
|
||||
frame = cv2.resize(center_crop_cv2(frame), (512, 512))
|
||||
pils.append(cv2_to_pil(frame))
|
||||
|
||||
return pils
|
||||
|
||||
def save_videos_from_pils(pils, path, fps=24):
|
||||
width, height = pils[0].size
|
||||
print(width, height)
|
||||
fourcc = cv2.VideoWriter_fourcc(*'mp4v')
|
||||
|
||||
pathlib.Path(path).parent.mkdir(exist_ok=True, parents=True)
|
||||
|
||||
output_name = pathlib.Path(path).stem
|
||||
temp_output_path = path.replace(output_name, output_name + '-temp')
|
||||
videowrite = cv2.VideoWriter(temp_output_path, fourcc, fps, (height, width))
|
||||
for pil in pils:
|
||||
frame = pil_to_cv2(pil)
|
||||
print(frame.shape, frame.min(), frame.max())
|
||||
videowrite.write(frame)
|
||||
|
||||
videowrite.release()
|
||||
"""
|
||||
embed()
|
||||
|
||||
cmd = (f'{get_ffmpeg_exe()} -i "{temp_output_path}"'
|
||||
f'-map 0:v -map 1:a -c:v h264 -shortest -y "{path}" -loglevel quiet')
|
||||
os.system(cmd)
|
||||
os.remove(temp_output_path)
|
||||
"""
|
||||
def save_video_from_cv2_list(pic_cv2_list, output_path, fps=30.0):
|
||||
pathlib.Path(output_path).parent.mkdir(exist_ok=True, parents=True)
|
||||
|
||||
num_frames = len(pic_cv2_list)
|
||||
height, width = pic_cv2_list[0].shape[:2]
|
||||
|
||||
#video_tensor = video_tensor[0, ...]
|
||||
#_, num_frames, height, width = video_tensor.shape
|
||||
|
||||
output_name = pathlib.Path(output_path).stem
|
||||
temp_output_path = output_path.replace(output_name, output_name)
|
||||
video_writer = cv2.VideoWriter(temp_output_path, cv2.VideoWriter_fourcc(*'mp4v'), fps, (width, height))
|
||||
|
||||
for i in range(num_frames):
|
||||
#frame_tensor = video_tensor[:, i, ...] # [c, h, w]
|
||||
#frame_tensor = frame_tensor.permute(1, 2, 0) # [h, w, c]
|
||||
|
||||
#frame_image = (frame_tensor * 255).numpy().astype(np.uint8)
|
||||
#frame_image = cv2.cvtColor(frame_image, cv2.COLOR_RGB2BGR)
|
||||
frame_image = pic_cv2_list[i].astype(np.uint8)
|
||||
video_writer.write(frame_image)
|
||||
|
||||
video_writer.release()
|
||||
|
||||
|
||||
#ffmpeg -i input_file -c:v libx264 -crf 20 -c:a aac -strict experimental -b:a 192k output_file
|
||||
@@ -0,0 +1,237 @@
|
||||
from skimage.transform import estimate_transform, AffineTransform
|
||||
import numpy as np
|
||||
from IPython import embed
|
||||
import mediapipe as mp
|
||||
import copy
|
||||
|
||||
mp_face_mesh = mp.solutions.face_mesh
|
||||
|
||||
|
||||
FACEMESH_LEFT_EYE = [i for i in mp_face_mesh.FACEMESH_LEFT_EYE]
|
||||
FACEMESH_RIGHT_EYE = [i for i in mp_face_mesh.FACEMESH_RIGHT_EYE]
|
||||
FACEMESH_LEFT_EYEBROW = [i for i in mp_face_mesh.FACEMESH_LEFT_EYEBROW]
|
||||
FACEMESH_RIGHT_EYEBROW = [i for i in mp_face_mesh.FACEMESH_RIGHT_EYEBROW]
|
||||
|
||||
# copy from draw_utils
|
||||
FACEMESH_LIPS_OUTER_BOTTOM_LEFT = [(61,146),(146,91),(91,181),(181,84),(84,17)]
|
||||
FACEMESH_LIPS_OUTER_BOTTOM_RIGHT = [(17,314),(314,405),(405,321),(321,375),(375,291)]
|
||||
FACEMESH_LIPS_INNER_BOTTOM_LEFT = [(78,95),(95,88),(88,178),(178,87),(87,14)]
|
||||
FACEMESH_LIPS_INNER_BOTTOM_RIGHT = [(14,317),(317,402),(402,318),(318,324),(324,308)]
|
||||
FACEMESH_LIPS_OUTER_TOP_LEFT = [(61,185),(185,40),(40,39),(39,37),(37,0)]
|
||||
FACEMESH_LIPS_OUTER_TOP_RIGHT = [(0,267),(267,269),(269,270),(270,409),(409,291)]
|
||||
FACEMESH_LIPS_INNER_TOP_LEFT = [(78,191),(191,80),(80,81),(81,82),(82,13)]
|
||||
FACEMESH_LIPS_INNER_TOP_RIGHT = [(13,312),(312,311),(311,310),(310,415),(415,308)]
|
||||
FACEMESH_MOUSE = \
|
||||
FACEMESH_LIPS_OUTER_BOTTOM_LEFT + \
|
||||
FACEMESH_LIPS_OUTER_BOTTOM_RIGHT + \
|
||||
FACEMESH_LIPS_INNER_BOTTOM_LEFT + \
|
||||
FACEMESH_LIPS_INNER_BOTTOM_RIGHT + \
|
||||
FACEMESH_LIPS_OUTER_TOP_LEFT + \
|
||||
FACEMESH_LIPS_OUTER_TOP_RIGHT + \
|
||||
FACEMESH_LIPS_INNER_TOP_LEFT + \
|
||||
FACEMESH_LIPS_INNER_TOP_RIGHT
|
||||
|
||||
LANDMARK_IDXES_DICT = {
|
||||
"left_eye" : sorted(list(set([j for i in FACEMESH_LEFT_EYE for j in i])) + [473]),
|
||||
"right_eye" : sorted(list(set([j for i in FACEMESH_RIGHT_EYE for j in i])) + [468]),
|
||||
"mouse" : sorted(list(set([j for i in FACEMESH_MOUSE for j in i]))),
|
||||
"nose" : sorted(list(set([1,4,5,274,275,281,44,45,51,220,440]))),
|
||||
"left_eyebow" : sorted(list(set([j for i in FACEMESH_LEFT_EYEBROW for j in i]))),
|
||||
"right_eyebow" : sorted(list(set([j for i in FACEMESH_RIGHT_EYEBROW for j in i]))),
|
||||
}
|
||||
|
||||
def create_perspective_matrix(aspect_ratio):
|
||||
kDegreesToRadians = np.pi / 180.
|
||||
near = 1
|
||||
far = 10000
|
||||
perspective_matrix = np.zeros(16, dtype=np.float32)
|
||||
|
||||
# Standard perspective projection matrix calculations.
|
||||
f = 1.0 / np.tan(kDegreesToRadians * 63 / 2.)
|
||||
|
||||
denom = 1.0 / (near - far)
|
||||
perspective_matrix[0] = f / aspect_ratio
|
||||
perspective_matrix[5] = f
|
||||
perspective_matrix[10] = (near + far) * denom
|
||||
perspective_matrix[11] = -1.
|
||||
perspective_matrix[14] = 1. * far * near * denom
|
||||
|
||||
# If the environment's origin point location is in the top left corner,
|
||||
# then skip additional flip along Y-axis is required to render correctly.
|
||||
|
||||
perspective_matrix[5] *= -1.
|
||||
return perspective_matrix
|
||||
|
||||
|
||||
def project_points_with_trans(points_3d, transformation_matrix, image_shape):
|
||||
P = create_perspective_matrix(image_shape[1] / image_shape[0]).reshape(4, 4).T
|
||||
L, N, _ = points_3d.shape
|
||||
projected_points = np.zeros((L, N, 2))
|
||||
#embed()
|
||||
for i in range(L):
|
||||
points_3d_frame = points_3d[i]
|
||||
ones = np.ones((points_3d_frame.shape[0], 1))
|
||||
points_3d_homogeneous = np.hstack([points_3d_frame, ones])
|
||||
transformed_points = points_3d_homogeneous @ transformation_matrix[i].T @ P
|
||||
projected_points_frame = transformed_points[:, :2] / transformed_points[:, 3, np.newaxis] # -1 ~ 1
|
||||
projected_points_frame[:, 0] = (projected_points_frame[:, 0] + 1) * 0.5 * image_shape[1]
|
||||
projected_points_frame[:, 1] = (projected_points_frame[:, 1] + 1) * 0.5 * image_shape[0]
|
||||
projected_points[i] = projected_points_frame
|
||||
return projected_points
|
||||
|
||||
def project_vertices_from_ref2tgt(ref_lmks3d, tgt_trans_mat):
|
||||
#eye_point_idxes
|
||||
projected_vertices = project_points_with_trans(ref_lmks3d[np.newaxis, ...], tgt_trans_mat[np.newaxis, ...], [512, 512])[0]
|
||||
return projected_vertices
|
||||
|
||||
|
||||
def old_motion_sync_old(sequence_driver_det, reference_det):
|
||||
assert type(sequence_driver_det) is list
|
||||
assert type(sequence_driver_det[0]) is type(reference_det)
|
||||
|
||||
lmks3d_mean = sum([i["lmks3d"] for i in sequence_driver_det]) / len(sequence_driver_det)
|
||||
overall_transform = estimate_transform('affine', lmks3d_mean, reference_det["lmks3d"])
|
||||
|
||||
eye_idxes_all = LANDMARK_IDXES_DICT["left_eye"] + LANDMARK_IDXES_DICT["right_eye"]
|
||||
for det_id in range(len(sequence_driver_det)):
|
||||
trans = estimate_transform('affine', sequence_driver_det[det_id]["lmks"][eye_idxes_all], sequence_driver_det[det_id]["lmks3d"][eye_idxes_all])
|
||||
sequence_driver_det[det_id]["lmks3d"] = np.vstack([
|
||||
sequence_driver_det[det_id]["lmks3d"],
|
||||
trans(sequence_driver_det[det_id]["lmks"][-10:])
|
||||
])
|
||||
|
||||
trans_mats = []
|
||||
for det in sequence_driver_det:
|
||||
trans_mats.append(det["trans_mat"] @ np.linalg.inv(sequence_driver_det[0]["trans_mat"]))
|
||||
|
||||
trans_mats_smooth = []
|
||||
smooth_margin = 2
|
||||
for tm_itx in range(len(trans_mats)):
|
||||
smooth_idxes = [i for i in range(tm_itx - smooth_margin, tm_itx + smooth_margin + 1) if i >= 0 and i < len(trans_mats)]
|
||||
tm = sum([trans_mats[i] for i in smooth_idxes]) / len(smooth_idxes)
|
||||
trans_mats_smooth.append(tm)
|
||||
|
||||
lmks3d_smooth = []
|
||||
smooth_margin = 1
|
||||
for det_itx in range(len(sequence_driver_det)):
|
||||
smooth_idxes = [i for i in range(det_itx - smooth_margin, det_itx + smooth_margin + 1) if i >= 0 and i < len(sequence_driver_det)]
|
||||
lmks3d_smooth.append(sum([sequence_driver_det[i]["lmks3d"] for i in smooth_idxes]) / len(smooth_idxes))
|
||||
|
||||
for det_itx, lmks3d in enumerate(lmks3d_smooth):
|
||||
sequence_driver_det[det_itx]["lmks3d"] = lmks3d
|
||||
|
||||
projected_vertices_list = []
|
||||
for det_itx in range(len(sequence_driver_det)):
|
||||
aligned_3d = overall_transform(sequence_driver_det[det_itx]["lmks3d"])
|
||||
tmat = reference_det["trans_mat"] @ trans_mats_smooth[det_itx]
|
||||
projected_vertices = project_vertices_from_ref2tgt(aligned_3d, tmat)
|
||||
projected_vertices_list.append(projected_vertices)
|
||||
|
||||
# note : use normed=False after motion_sync, when draw(ing)_landmarks
|
||||
# kps_image = vis.draw_landmarks((512, 512), projected_vertices, normed=False)
|
||||
return projected_vertices_list
|
||||
|
||||
|
||||
def motion_sync(sequence_driver_det, reference_det, per_landmark_align=True):
|
||||
assert type(sequence_driver_det) is list
|
||||
assert type(sequence_driver_det[0]) is type(reference_det)
|
||||
|
||||
eye_idxes_all = [i for i in sorted(list(set(LANDMARK_IDXES_DICT["left_eye"] + LANDMARK_IDXES_DICT["right_eye"]))) if i < len(reference_det["lmks3d"])]
|
||||
for det_id in range(len(sequence_driver_det)):
|
||||
trans_iris = estimate_transform('affine', sequence_driver_det[det_id]["lmks"][eye_idxes_all], sequence_driver_det[det_id]["lmks3d"][eye_idxes_all])
|
||||
sequence_driver_det[det_id]["lmks3d"] = np.vstack([
|
||||
sequence_driver_det[det_id]["lmks3d"],
|
||||
trans_iris(sequence_driver_det[det_id]["lmks"][-10:])
|
||||
])
|
||||
|
||||
trans_iris = estimate_transform('affine', reference_det["lmks"][eye_idxes_all], reference_det["lmks3d"][eye_idxes_all])
|
||||
reference_det["lmks3d"] = np.vstack([
|
||||
reference_det["lmks3d"],
|
||||
trans_iris(reference_det["lmks"][-10:])
|
||||
])
|
||||
|
||||
lmks3d_mean = sum([i["lmks3d"] for i in sequence_driver_det]) / len(sequence_driver_det)
|
||||
|
||||
landmark_trans_dict = {}
|
||||
for landmark_name, landmark_idxes in LANDMARK_IDXES_DICT.items():
|
||||
rf_lm = reference_det["lmks3d"][landmark_idxes]
|
||||
dr_lm = lmks3d_mean[landmark_idxes]
|
||||
landmark_trans_dict[landmark_name] = estimate_transform('affine', dr_lm, rf_lm)
|
||||
|
||||
#embed()
|
||||
overall_transform = estimate_transform('affine', lmks3d_mean, reference_det["lmks3d"])
|
||||
#embed()
|
||||
#lmks3d_mean = sum([i["lmks3d"] for i in sequence_driver_det]) / len(sequence_driver_det)
|
||||
#overall_transform = estimate_transform('affine', lmks3d_mean, reference_det["lmks3d"])
|
||||
|
||||
#driver_start_center = sequence_driver_det[0]["lmks3d"].mean(axis=0)
|
||||
#reference_center = reference_det["lmks3d"].mean(axis=0)
|
||||
#driver_start_size = ((sequence_driver_det[0]["lmks3d"] - driver_start_center)**2).sum()**(0.5)
|
||||
#reference_size = ((reference_det["lmks3d"] - reference_center)**2).sum()**(0.5)
|
||||
|
||||
#reference_det_lmks3d_rescale = (reference_det["lmks3d"] - reference_center) / reference_size * driver_start_size + driver_start_center
|
||||
#reference_transform_back = estimate_transform('affine', reference_det_lmks3d_rescale, reference_det["lmks3d"])
|
||||
|
||||
#driver_lmks3d_mean = sum([i["lmks3d"] for i in sequence_driver_det]) / len(sequence_driver_det)
|
||||
#facial_transform = estimate_transform('affine', driver_lmks3d_mean, reference_det_lmks3d_rescale)
|
||||
|
||||
|
||||
#for det_id in range(len(sequence_driver_det)):
|
||||
# trans = estimate_transform('affine', sequence_driver_det[det_id]["lmks"][:-10], sequence_driver_det[det_id]["lmks3d"])
|
||||
# sequence_driver_det[det_id]["lmks3d"] = trans(sequence_driver_det[det_id]["lmks"])
|
||||
|
||||
trans_mats = []
|
||||
for det in sequence_driver_det:
|
||||
trans_mats.append(det["trans_mat"] @ np.linalg.inv(sequence_driver_det[0]["trans_mat"]))
|
||||
|
||||
trans_mats_smooth = []
|
||||
smooth_margin = 2
|
||||
for tm_itx in range(len(trans_mats)):
|
||||
smooth_idxes = [i for i in range(tm_itx - smooth_margin, tm_itx + smooth_margin + 1) if i >= 0 and i < len(trans_mats)]
|
||||
tm = sum([trans_mats[i] for i in smooth_idxes]) / len(smooth_idxes)
|
||||
trans_mats_smooth.append(tm)
|
||||
|
||||
lmks3d_smooth = []
|
||||
smooth_margin = 1
|
||||
for det_itx in range(len(sequence_driver_det)):
|
||||
smooth_idxes = [i for i in range(det_itx - smooth_margin, det_itx + smooth_margin + 1) if i >= 0 and i < len(sequence_driver_det)]
|
||||
lmks3d_smooth.append(sum([sequence_driver_det[i]["lmks3d"] for i in smooth_idxes]) / len(smooth_idxes))
|
||||
|
||||
for det_itx, lmks3d in enumerate(lmks3d_smooth):
|
||||
sequence_driver_det[det_itx]["lmks3d"] = lmks3d
|
||||
|
||||
projected_vertices_list = []
|
||||
for det_itx in range(len(sequence_driver_det)):
|
||||
#aligned_3d = overall_transform(sequence_driver_det[det_itx]["lmks3d"])
|
||||
aligned_3d = copy.deepcopy(sequence_driver_det[det_itx]["lmks3d"])
|
||||
if per_landmark_align:
|
||||
for landmark_name, landmark_idxes in LANDMARK_IDXES_DICT.items():
|
||||
dr_lm = sequence_driver_det[det_itx]["lmks3d"][landmark_idxes]
|
||||
lm_trans = landmark_trans_dict[landmark_name]
|
||||
aligned_3d[landmark_idxes] = lm_trans(dr_lm)
|
||||
|
||||
#aligned_3d = lmks3d_mean
|
||||
tmat = trans_mats_smooth[det_itx] @ reference_det["trans_mat"]
|
||||
projected_vertices = project_vertices_from_ref2tgt(aligned_3d, tmat)
|
||||
projected_vertices_list.append(projected_vertices)
|
||||
|
||||
continue
|
||||
|
||||
trans_ref_aligned_to_driver = (sequence_driver_det[det_itx]["trans_mat"]) @ np.linalg.inv(reference_det["trans_mat"])
|
||||
ref_aligned_to_driver = AffineTransform(trans_ref_aligned_to_driver)(reference_det["lmks3d"])
|
||||
det["trans_mat"] @ np.linalg.inv(sequence_driver_det[0]["trans_mat"])
|
||||
aligned_3d = sequence_driver_det[det_itx]["lmks3d"]
|
||||
|
||||
|
||||
#facial_transform(sequence_driver_det[det_itx]["lmks3d"])
|
||||
#tmat = reference_det["trans_mat"] @ trans_mats_smooth[det_itx]
|
||||
|
||||
tmat = sequence_driver_det[det_itx]["trans_mat"] @ trans_mats_smooth[det_itx] #@ reference_transform_back.params
|
||||
projected_vertices = project_vertices_from_ref2tgt(aligned_3d, tmat)
|
||||
#embed()
|
||||
#projected_vertices = reference_transform_back(projected_vertices)
|
||||
projected_vertices_list.append(projected_vertices)
|
||||
|
||||
# note : use normed=False after motion_sync, when draw(ing)_landmarks
|
||||
# kps_image = vis.draw_landmarks((512, 512), projected_vertices, normed=False)
|
||||
return projected_vertices_list
|
||||
Binary file not shown.
Binary file not shown.
@@ -0,0 +1,94 @@
|
||||
import os
|
||||
import numpy as np
|
||||
import cv2
|
||||
import time
|
||||
from tqdm import tqdm
|
||||
import multiprocessing
|
||||
import glob
|
||||
|
||||
import mediapipe as mp
|
||||
from mediapipe import solutions
|
||||
from mediapipe.framework.formats import landmark_pb2
|
||||
from mediapipe.tasks import python
|
||||
from mediapipe.tasks.python import vision
|
||||
from . import face_landmark
|
||||
|
||||
CUR_DIR = os.path.dirname(__file__)
|
||||
|
||||
|
||||
class LMKExtractor():
|
||||
def __init__(self, FPS=25):
|
||||
# Create an FaceLandmarker object.
|
||||
self.mode = mp.tasks.vision.FaceDetectorOptions.running_mode.IMAGE
|
||||
base_options = python.BaseOptions(model_asset_path=os.path.join(CUR_DIR, 'mp_models/face_landmarker_v2_with_blendshapes.task'))
|
||||
base_options.delegate = mp.tasks.BaseOptions.Delegate.CPU
|
||||
options = vision.FaceLandmarkerOptions(base_options=base_options,
|
||||
running_mode=self.mode,
|
||||
output_face_blendshapes=True,
|
||||
output_facial_transformation_matrixes=True,
|
||||
num_faces=1)
|
||||
self.detector = face_landmark.FaceLandmarker.create_from_options(options)
|
||||
self.last_ts = 0
|
||||
self.frame_ms = int(1000 / FPS)
|
||||
|
||||
det_base_options = python.BaseOptions(model_asset_path=os.path.join(CUR_DIR, 'mp_models/blaze_face_short_range.tflite'))
|
||||
det_options = vision.FaceDetectorOptions(base_options=det_base_options)
|
||||
self.det_detector = vision.FaceDetector.create_from_options(det_options)
|
||||
|
||||
|
||||
def __call__(self, img):
|
||||
frame = cv2.cvtColor(img, cv2.COLOR_BGR2RGB)
|
||||
image = mp.Image(image_format=mp.ImageFormat.SRGB, data=frame)
|
||||
t0 = time.time()
|
||||
if self.mode == mp.tasks.vision.FaceDetectorOptions.running_mode.VIDEO:
|
||||
det_result = self.det_detector.detect(image)
|
||||
if len(det_result.detections) != 1:
|
||||
return None
|
||||
self.last_ts += self.frame_ms
|
||||
try:
|
||||
detection_result, mesh3d = self.detector.detect_for_video(image, timestamp_ms=self.last_ts)
|
||||
except:
|
||||
return None
|
||||
elif self.mode == mp.tasks.vision.FaceDetectorOptions.running_mode.IMAGE:
|
||||
# det_result = self.det_detector.detect(image)
|
||||
|
||||
# if len(det_result.detections) != 1:
|
||||
# return None
|
||||
try:
|
||||
detection_result, mesh3d = self.detector.detect(image)
|
||||
except:
|
||||
return None
|
||||
|
||||
|
||||
bs_list = detection_result.face_blendshapes
|
||||
if len(bs_list) == 1:
|
||||
bs = bs_list[0]
|
||||
bs_values = []
|
||||
for index in range(len(bs)):
|
||||
bs_values.append(bs[index].score)
|
||||
bs_values = bs_values[1:] # remove neutral
|
||||
trans_mat = detection_result.facial_transformation_matrixes[0]
|
||||
face_landmarks_list = detection_result.face_landmarks
|
||||
face_landmarks = face_landmarks_list[0]
|
||||
lmks = []
|
||||
for index in range(len(face_landmarks)):
|
||||
x = face_landmarks[index].x
|
||||
y = face_landmarks[index].y
|
||||
z = face_landmarks[index].z
|
||||
lmks.append([x, y, z])
|
||||
lmks = np.array(lmks)
|
||||
|
||||
lmks3d = np.array(mesh3d.vertex_buffer)
|
||||
lmks3d = lmks3d.reshape(-1, 5)[:, :3]
|
||||
mp_tris = np.array(mesh3d.index_buffer).reshape(-1, 3) + 1
|
||||
|
||||
return {
|
||||
"lmks": lmks,
|
||||
'lmks3d': lmks3d,
|
||||
"trans_mat": trans_mat,
|
||||
'faces': mp_tris,
|
||||
"bs": bs_values
|
||||
}
|
||||
else:
|
||||
# print('multiple faces in the image: {}'.format(img_path))
|
||||
return None
|
||||
@@ -0,0 +1,36 @@
|
||||
import torch
|
||||
import numpy as np
|
||||
|
||||
def get_alpha(alphas_cumprod, timestep):
|
||||
timestep_lt_zero_mask = torch.lt(timestep, 0).to(alphas_cumprod.dtype)
|
||||
normal_alpha = alphas_cumprod[torch.clip(timestep, 0)]
|
||||
one_alpha = torch.ones_like(normal_alpha).to(normal_alpha.dtype).to(normal_alpha.dtype)
|
||||
return normal_alpha * (1 - timestep_lt_zero_mask) + one_alpha * timestep_lt_zero_mask
|
||||
|
||||
def psuedo_velocity_wrt_noisy_and_timestep(noisy_images, noisy_images_pre, alphas_cumprod, timestep, timestep_prev):
|
||||
alpha_prod_t = get_alpha(alphas_cumprod, timestep).view(-1, 1, 1, 1, 1).detach()
|
||||
beta_prod_t = 1 - alpha_prod_t
|
||||
alpha_prod_t_prev = get_alpha(alphas_cumprod, timestep_prev).view(-1, 1, 1, 1, 1).detach()
|
||||
beta_prod_t_prev = 1 - alpha_prod_t_prev
|
||||
|
||||
a_s = (alpha_prod_t_prev ** (0.5)).to(noisy_images.dtype)
|
||||
a_t = (alpha_prod_t ** (0.5)).to(noisy_images.dtype)
|
||||
b_s = (beta_prod_t_prev ** (0.5)).to(noisy_images.dtype)
|
||||
b_t = (beta_prod_t ** (0.5)).to(noisy_images.dtype)
|
||||
|
||||
psuedo_velocity = (noisy_images_pre - (
|
||||
a_s * a_t + b_s * b_t
|
||||
) * noisy_images) / (
|
||||
b_s * a_t - a_s * b_t
|
||||
)
|
||||
|
||||
return psuedo_velocity
|
||||
|
||||
def origin_by_velocity_and_sample(velocity, noisy_images, alphas_cumprod, timestep):
|
||||
alpha_prod_t = get_alpha(alphas_cumprod, timestep).view(-1, 1, 1, 1, 1).detach()
|
||||
beta_prod_t = 1 - alpha_prod_t
|
||||
a_t = (alpha_prod_t ** (0.5)).to(noisy_images.dtype)
|
||||
b_t = (beta_prod_t ** (0.5)).to(noisy_images.dtype)
|
||||
|
||||
pred_original_sample = a_t * noisy_images - b_t * velocity
|
||||
return pred_original_sample
|
||||
@@ -0,0 +1,936 @@
|
||||
import importlib
|
||||
import os
|
||||
import os.path as osp
|
||||
import shutil
|
||||
import sys
|
||||
from pathlib import Path
|
||||
import av
|
||||
import numpy as np
|
||||
import torch
|
||||
import torchvision
|
||||
from einops import rearrange
|
||||
from PIL import Image
|
||||
import random
|
||||
import subprocess
|
||||
import sys
|
||||
from pathlib import Path
|
||||
from typing import List
|
||||
import cv2
|
||||
import mediapipe as mp
|
||||
from moviepy.editor import AudioFileClip, VideoClip
|
||||
|
||||
|
||||
def filter_non_none(dict_obj):
|
||||
"""
|
||||
Filters out key-value pairs from the given dictionary where the value is None.
|
||||
|
||||
Args:
|
||||
dict_obj (Dict): The dictionary to be filtered.
|
||||
|
||||
Returns:
|
||||
Dict: The dictionary with key-value pairs removed where the value was None.
|
||||
|
||||
This function creates a new dictionary containing only the key-value pairs from
|
||||
the original dictionary where the value is not None. It then clears the original
|
||||
dictionary and updates it with the filtered key-value pairs.
|
||||
"""
|
||||
non_none_filter = { k: v for k, v in dict_obj.items() if v is not None }
|
||||
dict_obj.clear()
|
||||
dict_obj.update(non_none_filter)
|
||||
return dict_obj
|
||||
|
||||
|
||||
def import_filename(filename):
|
||||
spec = importlib.util.spec_from_file_location("mymodule", filename)
|
||||
module = importlib.util.module_from_spec(spec)
|
||||
sys.modules[spec.name] = module
|
||||
spec.loader.exec_module(module)
|
||||
return module
|
||||
|
||||
|
||||
def save_videos_from_pil(pil_images, path, fps=8, audio_path=None):
|
||||
import av
|
||||
|
||||
save_fmt = Path(path).suffix
|
||||
os.makedirs(os.path.dirname(path), exist_ok=True)
|
||||
width, height = pil_images[0].size
|
||||
|
||||
if save_fmt == ".mp4":
|
||||
codec = "libx264"
|
||||
container = av.open(path, "w")
|
||||
stream = container.add_stream(codec, rate=fps)
|
||||
|
||||
stream.width = width
|
||||
stream.height = height
|
||||
|
||||
for pil_image in pil_images:
|
||||
# pil_image = Image.fromarray(image_arr).convert("RGB")
|
||||
av_frame = av.VideoFrame.from_image(pil_image)
|
||||
container.mux(stream.encode(av_frame))
|
||||
container.mux(stream.encode())
|
||||
container.close()
|
||||
|
||||
elif save_fmt == ".gif":
|
||||
pil_images[0].save(
|
||||
fp=path,
|
||||
format="GIF",
|
||||
append_images=pil_images[1:],
|
||||
save_all=True,
|
||||
duration=(1 / fps * 1000),
|
||||
loop=0,
|
||||
)
|
||||
else:
|
||||
raise ValueError("Unsupported file type. Use .mp4 or .gif.")
|
||||
|
||||
|
||||
def save_videos_grid(videos: torch.Tensor, path: str, audio_path=None, rescale=False, n_rows=6, fps=8):
|
||||
videos = rearrange(videos, "b c t h w -> t b c h w")
|
||||
height, width = videos.shape[-2:]
|
||||
outputs = []
|
||||
|
||||
for x in videos:
|
||||
x = torchvision.utils.make_grid(x, nrow=n_rows) # (c h w)
|
||||
x = x.transpose(0, 1).transpose(1, 2).squeeze(-1) # (h w c)
|
||||
if rescale:
|
||||
x = (x + 1.0) / 2.0 # -1,1 -> 0,1
|
||||
x = (x * 255).numpy().astype(np.uint8)
|
||||
x = Image.fromarray(x)
|
||||
|
||||
outputs.append(x)
|
||||
|
||||
os.makedirs(os.path.dirname(path), exist_ok=True)
|
||||
|
||||
save_videos_from_pil(outputs, path, fps, audio_path=audio_path)
|
||||
|
||||
|
||||
def read_frames(video_path):
|
||||
container = av.open(video_path)
|
||||
|
||||
video_stream = next(s for s in container.streams if s.type == "video")
|
||||
frames = []
|
||||
for packet in container.demux(video_stream):
|
||||
for frame in packet.decode():
|
||||
image = Image.frombytes(
|
||||
"RGB",
|
||||
(frame.width, frame.height),
|
||||
frame.to_rgb().to_ndarray(),
|
||||
)
|
||||
frames.append(image)
|
||||
|
||||
return frames
|
||||
|
||||
|
||||
def get_fps(video_path):
|
||||
container = av.open(video_path)
|
||||
video_stream = next(s for s in container.streams if s.type == "video")
|
||||
fps = video_stream.average_rate
|
||||
container.close()
|
||||
return fps
|
||||
|
||||
|
||||
def crop_and_pad(image, rect):
|
||||
x0, y0, x1, y1 = rect
|
||||
h, w = image.shape[:2]
|
||||
|
||||
# 确保坐标在图像范围内
|
||||
x0, y0 = max(0, x0), max(0, y0)
|
||||
x1, y1 = min(w, x1), min(h, y1)
|
||||
|
||||
# 计算原始框的宽度和高度
|
||||
width = x1 - x0
|
||||
height = y1 - y0
|
||||
|
||||
# 使用较小的边长作为裁剪正方形的边长
|
||||
side_length = min(width, height)
|
||||
|
||||
# 计算正方形框中心点
|
||||
center_x = (x0 + x1) // 2
|
||||
center_y = (y0 + y1) // 2
|
||||
|
||||
# 重新计算正方形框的坐标
|
||||
new_x0 = max(0, center_x - side_length // 2)
|
||||
new_y0 = max(0, center_y - side_length // 2)
|
||||
new_x1 = min(w, new_x0 + side_length)
|
||||
new_y1 = min(h, new_y0 + side_length)
|
||||
|
||||
# 最终裁剪框的尺寸修正（确保是正方形）
|
||||
if (new_x1 - new_x0) != (new_y1 - new_y0):
|
||||
side_length = min(new_x1 - new_x0, new_y1 - new_y0)
|
||||
new_x1 = new_x0 + side_length
|
||||
new_y1 = new_y0 + side_length
|
||||
|
||||
# 裁剪图像
|
||||
cropped_image = image[new_y0:new_y1, new_x0:new_x1]
|
||||
|
||||
return cropped_image, (new_x0, new_y0, new_x1, new_y1)
|
||||
|
||||
|
||||
|
||||
def seed_everything(seed):
|
||||
"""
|
||||
Seeds all random number generators to ensure reproducibility.
|
||||
|
||||
Args:
|
||||
seed (int): The seed value to set for all random number generators.
|
||||
"""
|
||||
torch.manual_seed(seed)
|
||||
torch.cuda.manual_seed_all(seed)
|
||||
np.random.seed(seed % (2**32))
|
||||
random.seed(seed)
|
||||
|
||||
|
||||
|
||||
def delete_additional_ckpt(base_path, num_keep):
|
||||
"""
|
||||
Deletes additional checkpoint files in the given directory.
|
||||
|
||||
Args:
|
||||
base_path (str): The path to the directory containing the checkpoint files.
|
||||
num_keep (int): The number of most recent checkpoint files to keep.
|
||||
|
||||
Returns:
|
||||
None
|
||||
|
||||
Raises:
|
||||
FileNotFoundError: If the base_path does not exist.
|
||||
|
||||
Example:
|
||||
>>> delete_additional_ckpt('path/to/checkpoints', 1)
|
||||
# This will delete all but the most recent checkpoint file in 'path/to/checkpoints'.
|
||||
"""
|
||||
dirs = []
|
||||
for d in os.listdir(base_path):
|
||||
if d.startswith("checkpoint-"):
|
||||
dirs.append(d)
|
||||
num_tot = len(dirs)
|
||||
if num_tot <= num_keep:
|
||||
return
|
||||
# ensure ckpt is sorted and delete the ealier!
|
||||
del_dirs = sorted(dirs, key=lambda x: int(
|
||||
x.split("-")[-1]))[: num_tot - num_keep]
|
||||
for d in del_dirs:
|
||||
path_to_dir = osp.join(base_path, d)
|
||||
if osp.exists(path_to_dir):
|
||||
shutil.rmtree(path_to_dir)
|
||||
|
||||
|
||||
|
||||
def tensor_to_video(tensor, output_video_file, audio_source, fps=25):
|
||||
"""
|
||||
Converts a Tensor with shape [c, f, h, w] into a video and adds an audio track from the specified audio file.
|
||||
|
||||
Args:
|
||||
tensor (Tensor): The Tensor to be converted, shaped [c, f, h, w].
|
||||
output_video_file (str): The file path where the output video will be saved.
|
||||
audio_source (str): The path to the audio file (WAV file) that contains the audio track to be added.
|
||||
fps (int): The frame rate of the output video. Default is 25 fps.
|
||||
"""
|
||||
tensor = tensor.permute(1, 2, 3, 0).cpu(
|
||||
).numpy() # convert to [f, h, w, c]
|
||||
tensor = np.clip(tensor * 255, 0, 255).astype(
|
||||
np.uint8
|
||||
) # to [0, 255]
|
||||
|
||||
def make_frame(t):
|
||||
# get index
|
||||
frame_index = min(int(t * fps), tensor.shape[0] - 1)
|
||||
return tensor[frame_index]
|
||||
new_video_clip = VideoClip(make_frame, duration=tensor.shape[0] / fps)
|
||||
audio_clip = AudioFileClip(audio_source).subclip(0, tensor.shape[0] / fps)
|
||||
new_video_clip = new_video_clip.set_audio(audio_clip)
|
||||
new_video_clip.write_videofile(output_video_file, fps=fps, audio_codec='aac')
|
||||
|
||||
|
||||
def tensor_to_video_new(tensor, output_video_file, audio_path, fps=25):
|
||||
|
||||
tensor = tensor.permute(1, 2, 3, 0).cpu(
|
||||
).numpy() # convert to [f, h, w, c]
|
||||
tensor = np.clip(tensor * 255, 0, 255).astype(
|
||||
np.uint8
|
||||
) # to [0, 255]
|
||||
|
||||
f, h, w, c = tensor.shape
|
||||
temp_video_path = output_video_file[:4] + "_no_audio.mp4"
|
||||
video_writer = cv2.VideoWriter(temp_video_path, cv2.VideoWriter_fourcc(*'mp4v'), 25, (w, h))
|
||||
|
||||
for t in tensor:
|
||||
t = cv2.cvtColor(t, cv2.COLOR_RGB2BGR)
|
||||
video_writer.write(t)
|
||||
video_writer.release()
|
||||
|
||||
cmd = (f'ffmpeg -i "{temp_video_path}" -i "{audio_path}" '
|
||||
f'-map 0:v -map 1:a -c:v h264 -shortest -y "{output_video_file}" -loglevel quiet')
|
||||
os.system(cmd)
|
||||
|
||||
|
||||
def save_video(frames_path, audio_path, output_path, frame_w, frame_h):
|
||||
|
||||
temp_video_path = output_path[:4] + "_no_audio.mp4"
|
||||
|
||||
command = f"ffmpeg -v error -hide_banner -framerate 25 -i {frames_path} -c:v libx264 -r 25 {temp_video_path} -y"
|
||||
subprocess.run(command, shell=True, check=True, text=True)
|
||||
|
||||
command = f"ffmpeg -i {temp_video_path} -i {audio_path} -map 0:v -map 1:a -c:v h264 -shortest -y {output_path} -loglevel quiet"
|
||||
subprocess.run(command, shell=True, check=True, text=True)
|
||||
|
||||
return
|
||||
|
||||
|
||||
silhouette_ids = [
|
||||
10, 338, 297, 332, 284, 251, 389, 356, 454, 323, 361, 288,
|
||||
397, 365, 379, 378, 400, 377, 152, 148, 176, 149, 150, 136,
|
||||
172, 58, 132, 93, 234, 127, 162, 21, 54, 103, 67, 109
|
||||
]
|
||||
lip_ids = [61, 185, 40, 39, 37, 0, 267, 269, 270, 409, 291,
|
||||
146, 91, 181, 84, 17, 314, 405, 321, 375]
|
||||
|
||||
|
||||
def compute_face_landmarks(detection_result, h, w):
|
||||
"""
|
||||
Compute face landmarks from a detection result.
|
||||
|
||||
Args:
|
||||
detection_result (mediapipe.solutions.face_mesh.FaceMesh): The detection result containing face landmarks.
|
||||
h (int): The height of the video frame.
|
||||
w (int): The width of the video frame.
|
||||
|
||||
Returns:
|
||||
face_landmarks_list (list): A list of face landmarks.
|
||||
"""
|
||||
face_landmarks_list = detection_result.face_landmarks
|
||||
if len(face_landmarks_list) != 1:
|
||||
print("#face is invalid:", len(face_landmarks_list))
|
||||
return []
|
||||
return [[p.x * w, p.y * h] for p in face_landmarks_list[0]]
|
||||
|
||||
|
||||
def get_landmark(file):
|
||||
"""
|
||||
This function takes a file as input and returns the facial landmarks detected in the file.
|
||||
|
||||
Args:
|
||||
file (str): The path to the file containing the video or image to be processed.
|
||||
|
||||
Returns:
|
||||
Tuple[List[float], List[float]]: A tuple containing two lists of floats representing the x and y coordinates of the facial landmarks.
|
||||
"""
|
||||
model_path = "pretrained_models/face_analysis/models/face_landmarker_v2_with_blendshapes.task"
|
||||
BaseOptions = mp.tasks.BaseOptions
|
||||
FaceLandmarker = mp.tasks.vision.FaceLandmarker
|
||||
FaceLandmarkerOptions = mp.tasks.vision.FaceLandmarkerOptions
|
||||
VisionRunningMode = mp.tasks.vision.RunningMode
|
||||
# Create a face landmarker instance with the video mode:
|
||||
options = FaceLandmarkerOptions(
|
||||
base_options=BaseOptions(model_asset_path=model_path),
|
||||
running_mode=VisionRunningMode.IMAGE,
|
||||
)
|
||||
|
||||
with FaceLandmarker.create_from_options(options) as landmarker:
|
||||
image = mp.Image.create_from_file(str(file))
|
||||
height, width = image.height, image.width
|
||||
face_landmarker_result = landmarker.detect(image)
|
||||
face_landmark = compute_face_landmarks(
|
||||
face_landmarker_result, height, width)
|
||||
|
||||
return np.array(face_landmark), height, width
|
||||
|
||||
|
||||
def get_landmark_overframes(landmark_model, frames_path):
|
||||
"""
|
||||
This function iterate frames and returns the facial landmarks detected in each frame.
|
||||
|
||||
Args:
|
||||
landmark_model: mediapipe landmark model instance
|
||||
frames_path (str): The path to the video frames.
|
||||
|
||||
Returns:
|
||||
List[List[float], float, float]: A List containing two lists of floats representing the x and y coordinates of the facial landmarks.
|
||||
"""
|
||||
|
||||
face_landmarks = []
|
||||
|
||||
for file in sorted(os.listdir(frames_path)):
|
||||
image = mp.Image.create_from_file(os.path.join(frames_path, file))
|
||||
height, width = image.height, image.width
|
||||
landmarker_result = landmark_model.detect(image)
|
||||
frame_landmark = compute_face_landmarks(
|
||||
landmarker_result, height, width)
|
||||
face_landmarks.append(frame_landmark)
|
||||
|
||||
return face_landmarks, height, width
|
||||
|
||||
|
||||
def get_lip_mask(landmarks, height, width, out_path=None, expand_ratio=2.0):
|
||||
"""
|
||||
Extracts the lip region from the given landmarks and saves it as an image.
|
||||
|
||||
Parameters:
|
||||
landmarks (numpy.ndarray): Array of facial landmarks.
|
||||
height (int): Height of the output lip mask image.
|
||||
width (int): Width of the output lip mask image.
|
||||
out_path (pathlib.Path): Path to save the lip mask image.
|
||||
expand_ratio (float): Expand ratio of mask.
|
||||
"""
|
||||
lip_landmarks = np.take(landmarks, lip_ids, 0)
|
||||
min_xy_lip = np.round(np.min(lip_landmarks, 0))
|
||||
max_xy_lip = np.round(np.max(lip_landmarks, 0))
|
||||
min_xy_lip[0], max_xy_lip[0], min_xy_lip[1], max_xy_lip[1] = expand_region(
|
||||
[min_xy_lip[0], max_xy_lip[0], min_xy_lip[1], max_xy_lip[1]], width, height, expand_ratio)
|
||||
lip_mask = np.zeros((height, width), dtype=np.uint8)
|
||||
lip_mask[round(min_xy_lip[1]):round(max_xy_lip[1]),
|
||||
round(min_xy_lip[0]):round(max_xy_lip[0])] = 255
|
||||
if out_path:
|
||||
cv2.imwrite(str(out_path), lip_mask)
|
||||
return None
|
||||
|
||||
return lip_mask
|
||||
|
||||
|
||||
def get_union_lip_mask(landmarks, height, width, expand_ratio=1):
|
||||
"""
|
||||
Extracts the lip region from the given landmarks and saves it as an image.
|
||||
|
||||
Parameters:
|
||||
landmarks (numpy.ndarray): Array of facial landmarks.
|
||||
height (int): Height of the output lip mask image.
|
||||
width (int): Width of the output lip mask image.
|
||||
expand_ratio (float): Expand ratio of mask.
|
||||
"""
|
||||
lip_masks = []
|
||||
for landmark in landmarks:
|
||||
lip_masks.append(get_lip_mask(landmarks=landmark, height=height,
|
||||
width=width, expand_ratio=expand_ratio))
|
||||
union_mask = get_union_mask(lip_masks)
|
||||
return union_mask
|
||||
|
||||
|
||||
def get_face_mask(landmarks, height, width, out_path=None, expand_ratio=1.2):
|
||||
"""
|
||||
Generate a face mask based on the given landmarks.
|
||||
|
||||
Args:
|
||||
landmarks (numpy.ndarray): The landmarks of the face.
|
||||
height (int): The height of the output face mask image.
|
||||
width (int): The width of the output face mask image.
|
||||
out_path (pathlib.Path): The path to save the face mask image.
|
||||
expand_ratio (float): Expand ratio of mask.
|
||||
Returns:
|
||||
None. The face mask image is saved at the specified path.
|
||||
"""
|
||||
face_landmarks = np.take(landmarks, silhouette_ids, 0)
|
||||
min_xy_face = np.round(np.min(face_landmarks, 0))
|
||||
max_xy_face = np.round(np.max(face_landmarks, 0))
|
||||
min_xy_face[0], max_xy_face[0], min_xy_face[1], max_xy_face[1] = expand_region(
|
||||
[min_xy_face[0], max_xy_face[0], min_xy_face[1], max_xy_face[1]], width, height, expand_ratio)
|
||||
face_mask = np.zeros((height, width), dtype=np.uint8)
|
||||
face_mask[round(min_xy_face[1]):round(max_xy_face[1]),
|
||||
round(min_xy_face[0]):round(max_xy_face[0])] = 255
|
||||
if out_path:
|
||||
cv2.imwrite(str(out_path), face_mask)
|
||||
return None
|
||||
|
||||
return face_mask
|
||||
|
||||
|
||||
def get_union_face_mask(landmarks, height, width, expand_ratio=1):
|
||||
"""
|
||||
Generate a face mask based on the given landmarks.
|
||||
|
||||
Args:
|
||||
landmarks (numpy.ndarray): The landmarks of the face.
|
||||
height (int): The height of the output face mask image.
|
||||
width (int): The width of the output face mask image.
|
||||
expand_ratio (float): Expand ratio of mask.
|
||||
Returns:
|
||||
None. The face mask image is saved at the specified path.
|
||||
"""
|
||||
face_masks = []
|
||||
for landmark in landmarks:
|
||||
face_masks.append(get_face_mask(landmarks=landmark,height=height,width=width,expand_ratio=expand_ratio))
|
||||
union_mask = get_union_mask(face_masks)
|
||||
return union_mask
|
||||
|
||||
def get_mask(file, cache_dir, face_expand_raio):
|
||||
"""
|
||||
Generate a face mask based on the given landmarks and save it to the specified cache directory.
|
||||
|
||||
Args:
|
||||
file (str): The path to the file containing the landmarks.
|
||||
cache_dir (str): The directory to save the generated face mask.
|
||||
|
||||
Returns:
|
||||
None
|
||||
"""
|
||||
landmarks, height, width = get_landmark(file)
|
||||
file_name = os.path.basename(file).split(".")[0]
|
||||
get_lip_mask(landmarks, height, width, os.path.join(
|
||||
cache_dir, f"{file_name}_lip_mask.png"))
|
||||
get_face_mask(landmarks, height, width, os.path.join(
|
||||
cache_dir, f"{file_name}_face_mask.png"), face_expand_raio)
|
||||
get_blur_mask(os.path.join(
|
||||
cache_dir, f"{file_name}_face_mask.png"), os.path.join(
|
||||
cache_dir, f"{file_name}_face_mask_blur.png"), kernel_size=(51, 51))
|
||||
get_blur_mask(os.path.join(
|
||||
cache_dir, f"{file_name}_lip_mask.png"), os.path.join(
|
||||
cache_dir, f"{file_name}_sep_lip.png"), kernel_size=(31, 31))
|
||||
get_background_mask(os.path.join(
|
||||
cache_dir, f"{file_name}_face_mask_blur.png"), os.path.join(
|
||||
cache_dir, f"{file_name}_sep_background.png"))
|
||||
get_sep_face_mask(os.path.join(
|
||||
cache_dir, f"{file_name}_face_mask_blur.png"), os.path.join(
|
||||
cache_dir, f"{file_name}_sep_lip.png"), os.path.join(
|
||||
cache_dir, f"{file_name}_sep_face.png"))
|
||||
|
||||
|
||||
def expand_region(region, image_w, image_h, expand_ratio=1.0):
|
||||
"""
|
||||
Expand the given region by a specified ratio.
|
||||
Args:
|
||||
region (tuple): A tuple containing the coordinates (min_x, max_x, min_y, max_y) of the region.
|
||||
image_w (int): The width of the image.
|
||||
image_h (int): The height of the image.
|
||||
expand_ratio (float, optional): The ratio by which the region should be expanded. Defaults to 1.0.
|
||||
|
||||
Returns:
|
||||
tuple: A tuple containing the expanded coordinates (min_x, max_x, min_y, max_y) of the region.
|
||||
"""
|
||||
|
||||
min_x, max_x, min_y, max_y = region
|
||||
mid_x = (max_x + min_x) // 2
|
||||
side_len_x = (max_x - min_x) * expand_ratio
|
||||
mid_y = (max_y + min_y) // 2
|
||||
side_len_y = (max_y - min_y) * expand_ratio
|
||||
min_x = mid_x - side_len_x // 2
|
||||
max_x = mid_x + side_len_x // 2
|
||||
min_y = mid_y - side_len_y // 2
|
||||
max_y = mid_y + side_len_y // 2
|
||||
if min_x < 0:
|
||||
max_x -= min_x
|
||||
min_x = 0
|
||||
if max_x > image_w:
|
||||
min_x -= max_x - image_w
|
||||
max_x = image_w
|
||||
if min_y < 0:
|
||||
max_y -= min_y
|
||||
min_y = 0
|
||||
if max_y > image_h:
|
||||
min_y -= max_y - image_h
|
||||
max_y = image_h
|
||||
|
||||
return round(min_x), round(max_x), round(min_y), round(max_y)
|
||||
|
||||
|
||||
def get_blur_mask(file_path, output_file_path, resize_dim=(64, 64), kernel_size=(101, 101)):
|
||||
"""
|
||||
Read, resize, blur, normalize, and save an image.
|
||||
|
||||
Parameters:
|
||||
file_path (str): Path to the input image file.
|
||||
output_dir (str): Path to the output directory to save blurred images.
|
||||
resize_dim (tuple): Dimensions to resize the images to.
|
||||
kernel_size (tuple): Size of the kernel to use for Gaussian blur.
|
||||
"""
|
||||
# Read the mask image
|
||||
mask = cv2.imread(file_path, cv2.IMREAD_GRAYSCALE)
|
||||
|
||||
# Check if the image is loaded successfully
|
||||
if mask is not None:
|
||||
normalized_mask = blur_mask(mask,resize_dim=resize_dim,kernel_size=kernel_size)
|
||||
# Save the normalized mask image
|
||||
cv2.imwrite(output_file_path, normalized_mask)
|
||||
return f"Processed, normalized, and saved: {output_file_path}"
|
||||
return f"Failed to load image: {file_path}"
|
||||
|
||||
|
||||
def blur_mask(mask, resize_dim=(64, 64), kernel_size=(51, 51)):
|
||||
"""
|
||||
Read, resize, blur, normalize, and save an image.
|
||||
|
||||
Parameters:
|
||||
file_path (str): Path to the input image file.
|
||||
resize_dim (tuple): Dimensions to resize the images to.
|
||||
kernel_size (tuple): Size of the kernel to use for Gaussian blur.
|
||||
"""
|
||||
# Check if the image is loaded successfully
|
||||
normalized_mask = None
|
||||
if mask is not None:
|
||||
# Resize the mask image
|
||||
resized_mask = cv2.resize(mask, resize_dim)
|
||||
# Apply Gaussian blur to the resized mask image
|
||||
blurred_mask = cv2.GaussianBlur(resized_mask, kernel_size, 0)
|
||||
# Normalize the blurred image
|
||||
normalized_mask = cv2.normalize(
|
||||
blurred_mask, None, 0, 255, cv2.NORM_MINMAX)
|
||||
# Save the normalized mask image
|
||||
return normalized_mask
|
||||
|
||||
def get_background_mask(file_path, output_file_path):
|
||||
"""
|
||||
Read an image, invert its values, and save the result.
|
||||
|
||||
Parameters:
|
||||
file_path (str): Path to the input image file.
|
||||
output_dir (str): Path to the output directory to save the inverted image.
|
||||
"""
|
||||
# Read the image
|
||||
image = cv2.imread(file_path, cv2.IMREAD_GRAYSCALE)
|
||||
|
||||
if image is None:
|
||||
print(f"Failed to load image: {file_path}")
|
||||
return
|
||||
|
||||
# Invert the image
|
||||
inverted_image = 1.0 - (
|
||||
image / 255.0
|
||||
) # Assuming the image values are in [0, 255] range
|
||||
# Convert back to uint8
|
||||
inverted_image = (inverted_image * 255).astype(np.uint8)
|
||||
|
||||
# Save the inverted image
|
||||
cv2.imwrite(output_file_path, inverted_image)
|
||||
# print(f"Processed and saved: {output_file_path}")
|
||||
|
||||
|
||||
def get_sep_face_mask(file_path1, file_path2, output_file_path):
|
||||
"""
|
||||
Read two images, subtract the second one from the first, and save the result.
|
||||
|
||||
Parameters:
|
||||
output_dir (str): Path to the output directory to save the subtracted image.
|
||||
"""
|
||||
|
||||
# Read the images
|
||||
mask1 = cv2.imread(file_path1, cv2.IMREAD_GRAYSCALE)
|
||||
mask2 = cv2.imread(file_path2, cv2.IMREAD_GRAYSCALE)
|
||||
|
||||
if mask1 is None or mask2 is None:
|
||||
print(f"Failed to load images: {file_path1}")
|
||||
return
|
||||
|
||||
# Ensure the images are the same size
|
||||
if mask1.shape != mask2.shape:
|
||||
print(
|
||||
f"Image shapes do not match for {file_path1}: {mask1.shape} vs {mask2.shape}"
|
||||
)
|
||||
return
|
||||
|
||||
# Subtract the second mask from the first
|
||||
result_mask = cv2.subtract(mask1, mask2)
|
||||
|
||||
# Save the result mask image
|
||||
cv2.imwrite(output_file_path, result_mask)
|
||||
# print(f"Processed and saved: {output_file_path}")
|
||||
|
||||
def resample_audio(input_audio_file: str, output_audio_file: str, sample_rate: int):
|
||||
p = subprocess.Popen([
|
||||
"ffmpeg", "-y", "-v", "error", "-i", input_audio_file, "-ar", str(sample_rate), output_audio_file
|
||||
])
|
||||
ret = p.wait()
|
||||
assert ret == 0, "Resample audio failed!"
|
||||
return output_audio_file
|
||||
|
||||
def get_face_region(image_path: str, detector):
|
||||
try:
|
||||
image = cv2.imread(image_path)
|
||||
if image is None:
|
||||
print(f"Failed to open image: {image_path}. Skipping...")
|
||||
return None, None
|
||||
|
||||
mp_image = mp.Image(image_format=mp.ImageFormat.SRGB, data=image)
|
||||
detection_result = detector.detect(mp_image)
|
||||
|
||||
# Adjust mask creation for the three-channel image
|
||||
mask = np.zeros_like(image, dtype=np.uint8)
|
||||
|
||||
for detection in detection_result.detections:
|
||||
bbox = detection.bounding_box
|
||||
start_point = (int(bbox.origin_x), int(bbox.origin_y))
|
||||
end_point = (int(bbox.origin_x + bbox.width),
|
||||
int(bbox.origin_y + bbox.height))
|
||||
cv2.rectangle(mask, start_point, end_point,
|
||||
(255, 255, 255), thickness=-1)
|
||||
|
||||
save_path = image_path.replace("images", "face_masks")
|
||||
os.makedirs(os.path.dirname(save_path), exist_ok=True)
|
||||
cv2.imwrite(save_path, mask)
|
||||
# print(f"Processed and saved {save_path}")
|
||||
return image_path, mask
|
||||
except Exception as e:
|
||||
print(f"Error processing image {image_path}: {e}")
|
||||
return None, None
|
||||
|
||||
|
||||
def save_checkpoint(model: torch.nn.Module, save_dir: str, prefix: str, ckpt_num: int, total_limit: int = -1) -> None:
|
||||
"""
|
||||
Save the model's state_dict to a checkpoint file.
|
||||
|
||||
If `total_limit` is provided, this function will remove the oldest checkpoints
|
||||
until the total number of checkpoints is less than the specified limit.
|
||||
|
||||
Args:
|
||||
model (nn.Module): The model whose state_dict is to be saved.
|
||||
save_dir (str): The directory where the checkpoint will be saved.
|
||||
prefix (str): The prefix for the checkpoint file name.
|
||||
ckpt_num (int): The checkpoint number to be saved.
|
||||
total_limit (int, optional): The maximum number of checkpoints to keep.
|
||||
Defaults to None, in which case no checkpoints will be removed.
|
||||
|
||||
Raises:
|
||||
FileNotFoundError: If the save directory does not exist.
|
||||
ValueError: If the checkpoint number is negative.
|
||||
OSError: If there is an error saving the checkpoint.
|
||||
"""
|
||||
|
||||
if not osp.exists(save_dir):
|
||||
raise FileNotFoundError(
|
||||
f"The save directory {save_dir} does not exist.")
|
||||
|
||||
if ckpt_num < 0:
|
||||
raise ValueError(f"Checkpoint number {ckpt_num} must be non-negative.")
|
||||
|
||||
save_path = osp.join(save_dir, f"{prefix}-{ckpt_num}.pth")
|
||||
|
||||
if total_limit > 0:
|
||||
checkpoints = os.listdir(save_dir)
|
||||
checkpoints = [d for d in checkpoints if d.startswith(prefix)]
|
||||
checkpoints = sorted(
|
||||
checkpoints, key=lambda x: int(x.split("-")[1].split(".")[0])
|
||||
)
|
||||
|
||||
if len(checkpoints) >= total_limit:
|
||||
num_to_remove = len(checkpoints) - total_limit + 1
|
||||
removing_checkpoints = checkpoints[0:num_to_remove]
|
||||
print(
|
||||
f"{len(checkpoints)} checkpoints already exist, removing {len(removing_checkpoints)} checkpoints"
|
||||
)
|
||||
print(
|
||||
f"Removing checkpoints: {', '.join(removing_checkpoints)}"
|
||||
)
|
||||
|
||||
for removing_checkpoint in removing_checkpoints:
|
||||
removing_checkpoint_path = osp.join(
|
||||
save_dir, removing_checkpoint)
|
||||
try:
|
||||
os.remove(removing_checkpoint_path)
|
||||
except OSError as e:
|
||||
print(
|
||||
f"Error removing checkpoint {removing_checkpoint_path}: {e}")
|
||||
|
||||
state_dict = model.state_dict()
|
||||
try:
|
||||
torch.save(state_dict, save_path)
|
||||
print(f"Checkpoint saved at {save_path}")
|
||||
except OSError as e:
|
||||
raise OSError(f"Error saving checkpoint at {save_path}: {e}") from e
|
||||
|
||||
|
||||
def init_output_dir(dir_list: List[str]):
|
||||
"""
|
||||
Initialize the output directories.
|
||||
|
||||
This function creates the directories specified in the `dir_list`. If a directory already exists, it does nothing.
|
||||
|
||||
Args:
|
||||
dir_list (List[str]): List of directory paths to create.
|
||||
"""
|
||||
for path in dir_list:
|
||||
os.makedirs(path, exist_ok=True)
|
||||
|
||||
|
||||
def load_checkpoint(cfg, save_dir, accelerator):
|
||||
"""
|
||||
Load the most recent checkpoint from the specified directory.
|
||||
|
||||
This function loads the latest checkpoint from the `save_dir` if the `resume_from_checkpoint` parameter is set to "latest".
|
||||
If a specific checkpoint is provided in `resume_from_checkpoint`, it loads that checkpoint. If no checkpoint is found,
|
||||
it starts training from scratch.
|
||||
|
||||
Args:
|
||||
cfg: The configuration object containing training parameters.
|
||||
save_dir (str): The directory where checkpoints are saved.
|
||||
accelerator: The accelerator object for distributed training.
|
||||
|
||||
Returns:
|
||||
int: The global step at which to resume training.
|
||||
"""
|
||||
if cfg.resume_from_checkpoint != "latest":
|
||||
resume_dir = cfg.resume_from_checkpoint
|
||||
else:
|
||||
resume_dir = save_dir
|
||||
# Get the most recent checkpoint
|
||||
dirs = os.listdir(resume_dir)
|
||||
|
||||
dirs = [d for d in dirs if d.startswith("checkpoint")]
|
||||
if len(dirs) > 0:
|
||||
dirs = sorted(dirs, key=lambda x: int(x.split("-")[1]))
|
||||
path = dirs[-1]
|
||||
accelerator.load_state(os.path.join(resume_dir, path))
|
||||
accelerator.print(f"Resuming from checkpoint {path}")
|
||||
global_step = int(path.split("-")[1])
|
||||
else:
|
||||
accelerator.print(
|
||||
f"Could not find checkpoint under {resume_dir}, start training from scratch")
|
||||
global_step = 0
|
||||
|
||||
return global_step
|
||||
|
||||
|
||||
def compute_snr(noise_scheduler, timesteps):
|
||||
"""
|
||||
Computes SNR as per
|
||||
https://github.com/TiankaiHang/Min-SNR-Diffusion-Training/blob/
|
||||
521b624bd70c67cee4bdf49225915f5945a872e3/guided_diffusion/gaussian_diffusion.py#L847-L849
|
||||
"""
|
||||
alphas_cumprod = noise_scheduler.alphas_cumprod
|
||||
sqrt_alphas_cumprod = alphas_cumprod**0.5
|
||||
sqrt_one_minus_alphas_cumprod = (1.0 - alphas_cumprod) ** 0.5
|
||||
|
||||
# Expand the tensors.
|
||||
# Adapted from https://github.com/TiankaiHang/Min-SNR-Diffusion-Training/blob/
|
||||
# 521b624bd70c67cee4bdf49225915f5945a872e3/guided_diffusion/gaussian_diffusion.py#L1026
|
||||
sqrt_alphas_cumprod = sqrt_alphas_cumprod.to(device=timesteps.device)[
|
||||
timesteps
|
||||
].float()
|
||||
while len(sqrt_alphas_cumprod.shape) < len(timesteps.shape):
|
||||
sqrt_alphas_cumprod = sqrt_alphas_cumprod[..., None]
|
||||
alpha = sqrt_alphas_cumprod.expand(timesteps.shape)
|
||||
|
||||
sqrt_one_minus_alphas_cumprod = sqrt_one_minus_alphas_cumprod.to(
|
||||
device=timesteps.device
|
||||
)[timesteps].float()
|
||||
while len(sqrt_one_minus_alphas_cumprod.shape) < len(timesteps.shape):
|
||||
sqrt_one_minus_alphas_cumprod = sqrt_one_minus_alphas_cumprod[..., None]
|
||||
sigma = sqrt_one_minus_alphas_cumprod.expand(timesteps.shape)
|
||||
|
||||
# Compute SNR.
|
||||
snr = (alpha / sigma) ** 2
|
||||
return snr
|
||||
|
||||
|
||||
def extract_audio_from_videos(video_path: Path, audio_output_path: Path) -> Path:
|
||||
"""
|
||||
Extract audio from a video file and save it as a WAV file.
|
||||
|
||||
This function uses ffmpeg to extract the audio stream from a given video file and saves it as a WAV file
|
||||
in the specified output directory.
|
||||
|
||||
Args:
|
||||
video_path (Path): The path to the input video file.
|
||||
output_dir (Path): The directory where the extracted audio file will be saved.
|
||||
|
||||
Returns:
|
||||
Path: The path to the extracted audio file.
|
||||
|
||||
Raises:
|
||||
subprocess.CalledProcessError: If the ffmpeg command fails to execute.
|
||||
"""
|
||||
ffmpeg_command = [
|
||||
'ffmpeg', '-y',
|
||||
'-i', str(video_path),
|
||||
'-vn', '-acodec',
|
||||
"pcm_s16le", '-ar', '16000', '-ac', '2',
|
||||
str(audio_output_path)
|
||||
]
|
||||
|
||||
try:
|
||||
# print(f"Running command: {' '.join(ffmpeg_command)}")
|
||||
subprocess.run(ffmpeg_command)
|
||||
except subprocess.CalledProcessError as e:
|
||||
print(f"Error extracting audio from video: {e}")
|
||||
raise
|
||||
|
||||
return audio_output_path
|
||||
|
||||
|
||||
def convert_video_to_images(video_path: Path, output_dir: Path) -> Path:
|
||||
"""
|
||||
Convert a video file into a sequence of images.
|
||||
|
||||
This function uses ffmpeg to convert each frame of the given video file into an image. The images are saved
|
||||
in a directory named after the video file stem under the specified output directory.
|
||||
|
||||
Args:
|
||||
video_path (Path): The path to the input video file.
|
||||
output_dir (Path): The directory where the extracted images will be saved.
|
||||
|
||||
Returns:
|
||||
Path: The path to the directory containing the extracted images.
|
||||
|
||||
Raises:
|
||||
subprocess.CalledProcessError: If the ffmpeg command fails to execute.
|
||||
"""
|
||||
ffmpeg_command = [
|
||||
'ffmpeg',
|
||||
'-i', str(video_path),
|
||||
'-vf', 'fps=25',
|
||||
str(output_dir / '%04d.png')
|
||||
]
|
||||
|
||||
try:
|
||||
# print(f"Running command: {' '.join(ffmpeg_command)}")
|
||||
subprocess.run(ffmpeg_command)
|
||||
except subprocess.CalledProcessError as e:
|
||||
print(f"Error converting video to images: {e}")
|
||||
raise
|
||||
|
||||
return output_dir
|
||||
|
||||
|
||||
def get_union_mask(masks):
|
||||
"""
|
||||
Compute the union of a list of masks.
|
||||
|
||||
This function takes a list of masks and computes their union by taking the maximum value at each pixel location.
|
||||
Additionally, it finds the bounding box of the non-zero regions in the mask and sets the bounding box area to white.
|
||||
|
||||
Args:
|
||||
masks (list of np.ndarray): List of masks to be combined.
|
||||
|
||||
Returns:
|
||||
np.ndarray: The union of the input masks.
|
||||
"""
|
||||
union_mask = None
|
||||
for mask in masks:
|
||||
if union_mask is None:
|
||||
union_mask = mask
|
||||
else:
|
||||
union_mask = np.maximum(union_mask, mask)
|
||||
|
||||
if union_mask is not None:
|
||||
# Find the bounding box of the non-zero regions in the mask
|
||||
rows = np.any(union_mask, axis=1)
|
||||
cols = np.any(union_mask, axis=0)
|
||||
try:
|
||||
ymin, ymax = np.where(rows)[0][[0, -1]]
|
||||
xmin, xmax = np.where(cols)[0][[0, -1]]
|
||||
except Exception as e:
|
||||
print(str(e))
|
||||
return 0.0
|
||||
|
||||
# Set bounding box area to white
|
||||
union_mask[ymin: ymax + 1, xmin: xmax + 1] = np.max(union_mask)
|
||||
|
||||
return union_mask
|
||||
|
||||
|
||||
def move_final_checkpoint(save_dir, module_dir, prefix):
|
||||
"""
|
||||
Move the final checkpoint file to the save directory.
|
||||
|
||||
This function identifies the latest checkpoint file based on the given prefix and moves it to the specified save directory.
|
||||
|
||||
Args:
|
||||
save_dir (str): The directory where the final checkpoint file should be saved.
|
||||
module_dir (str): The directory containing the checkpoint files.
|
||||
prefix (str): The prefix used to identify checkpoint files.
|
||||
|
||||
Raises:
|
||||
ValueError: If no checkpoint files are found with the specified prefix.
|
||||
"""
|
||||
checkpoints = os.listdir(module_dir)
|
||||
checkpoints = [d for d in checkpoints if d.startswith(prefix)]
|
||||
checkpoints = sorted(
|
||||
checkpoints, key=lambda x: int(x.split("-")[1].split(".")[0])
|
||||
)
|
||||
shutil.copy2(os.path.join(
|
||||
module_dir, checkpoints[-1]), os.path.join(save_dir, prefix + '.pth'))
|
||||
@@ -0,0 +1,731 @@
|
||||
# pylint: disable=E1101,C0415,W0718,R0801
|
||||
# scripts/train_stage1.py
|
||||
"""
|
||||
This is the main training script for stage 1 of the project.
|
||||
It imports necessary packages, defines necessary classes and functions, and trains the model using the provided configuration.
|
||||
|
||||
The script includes the following classes and functions:
|
||||
|
||||
1. Net: A PyTorch model that takes noisy latents, timesteps, reference image latents, face embeddings,
|
||||
and face masks as input and returns the denoised latents.
|
||||
3. log_validation: A function that logs the validation information using the given VAE, image encoder,
|
||||
network, scheduler, accelerator, width, height, and configuration.
|
||||
4. train_stage1_process: A function that processes the training stage 1 using the given configuration.
|
||||
|
||||
The script also includes the necessary imports and a brief description of the purpose of the file.
|
||||
"""
|
||||
|
||||
import argparse
|
||||
import copy
|
||||
import logging
|
||||
import math
|
||||
import os
|
||||
import warnings
|
||||
from datetime import datetime
|
||||
|
||||
import cv2
|
||||
import diffusers
|
||||
import mlflow
|
||||
import numpy as np
|
||||
import torch
|
||||
import torch.nn.functional as F
|
||||
import torch.utils.checkpoint
|
||||
import transformers
|
||||
from accelerate import Accelerator
|
||||
from accelerate.logging import get_logger
|
||||
from accelerate.utils import DistributedDataParallelKwargs
|
||||
from diffusers import AutoencoderKL, DDIMScheduler
|
||||
from diffusers.optimization import get_scheduler
|
||||
from diffusers.utils import check_min_version
|
||||
from diffusers.utils.import_utils import is_xformers_available
|
||||
from omegaconf import OmegaConf
|
||||
from PIL import Image
|
||||
from torch import nn
|
||||
from tqdm.auto import tqdm
|
||||
from torch.utils.tensorboard import SummaryWriter
|
||||
|
||||
from .src.pipelines.face_animate_static import StaticPipeline
|
||||
from .src.datasets.mask_image import FaceMaskDataset, FaceMaskDatasetVal
|
||||
from .src.models.mutual_self_attention import ReferenceAttentionControl
|
||||
from .src.models.unet_2d_condition import UNet2DConditionModel
|
||||
from .src.models.unet_3d_echo import EchoUNet3DConditionModel
|
||||
from .src.utils.util import (compute_snr, delete_additional_ckpt,
|
||||
import_filename, init_output_dir,
|
||||
load_checkpoint, move_final_checkpoint,
|
||||
save_checkpoint, seed_everything)
|
||||
|
||||
warnings.filterwarnings("ignore")
|
||||
|
||||
# Will error if the minimal version of diffusers is not installed. Remove at your own risks.
|
||||
check_min_version("0.10.0.dev0")
|
||||
|
||||
logger = get_logger(__name__, log_level="INFO")
|
||||
|
||||
writer = SummaryWriter()
|
||||
|
||||
class Net(nn.Module):
|
||||
"""
|
||||
The Net class defines a neural network model that combines a reference UNet2DConditionModel,
|
||||
a denoising UNet3DConditionModel, a face locator, and other components to animate a face in a static image.
|
||||
|
||||
Args:
|
||||
reference_unet (UNet2DConditionModel): The reference UNet2DConditionModel used for face animation.
|
||||
denoising_unet (UNet3DConditionModel): The denoising UNet3DConditionModel used for face animation.
|
||||
face_locator (FaceLocator): The face locator model used for face animation.
|
||||
reference_control_writer: The reference control writer component.
|
||||
reference_control_reader: The reference control reader component.
|
||||
imageproj: The image projection model.
|
||||
|
||||
Forward method:
|
||||
noisy_latents (torch.Tensor): The noisy latents tensor.
|
||||
timesteps (torch.Tensor): The timesteps tensor.
|
||||
ref_image_latents (torch.Tensor): The reference image latents tensor.
|
||||
face_emb (torch.Tensor): The face embeddings tensor.
|
||||
face_mask (torch.Tensor): The face mask tensor.
|
||||
uncond_fwd (bool): A flag indicating whether to perform unconditional forward pass.
|
||||
|
||||
Returns:
|
||||
torch.Tensor: The output tensor of the neural network model.
|
||||
"""
|
||||
|
||||
def __init__(
|
||||
self,
|
||||
reference_unet: UNet2DConditionModel,
|
||||
denoising_unet: EchoUNet3DConditionModel,
|
||||
reference_control_writer: ReferenceAttentionControl,
|
||||
reference_control_reader: ReferenceAttentionControl,
|
||||
):
|
||||
super().__init__()
|
||||
self.reference_unet = reference_unet
|
||||
self.denoising_unet = denoising_unet
|
||||
self.reference_control_writer = reference_control_writer
|
||||
self.reference_control_reader = reference_control_reader
|
||||
|
||||
def forward(
|
||||
self,
|
||||
noisy_latents,
|
||||
timesteps,
|
||||
ref_image_latents,
|
||||
audio_emb,
|
||||
uncond_fwd: bool = False,
|
||||
):
|
||||
"""
|
||||
Forward pass of the model.
|
||||
Args:
|
||||
self (Net): The model instance.
|
||||
noisy_latents (torch.Tensor): Noisy latents.
|
||||
timesteps (torch.Tensor): Timesteps.
|
||||
ref_image_latents (torch.Tensor): Reference image latents.
|
||||
face_emb (torch.Tensor): Face embedding.
|
||||
face_mask (torch.Tensor): Face mask.
|
||||
uncond_fwd (bool, optional): Unconditional forward pass. Defaults to False.
|
||||
|
||||
Returns:
|
||||
torch.Tensor: Model prediction.
|
||||
"""
|
||||
|
||||
# audio_emb = torch.zeros((noisy_latents.shape[0],50,384)).to(
|
||||
# device=self.reference_unet.device, dtype=self.reference_unet.dtype)
|
||||
audio_emb = audio_emb.to(device=self.reference_unet.device, dtype=self.reference_unet.dtype)
|
||||
|
||||
if not uncond_fwd:
|
||||
ref_timesteps = torch.zeros_like(timesteps)
|
||||
self.reference_unet(
|
||||
ref_image_latents,
|
||||
ref_timesteps,
|
||||
encoder_hidden_states=None,
|
||||
return_dict=False,
|
||||
)
|
||||
self.reference_control_reader.update(self.reference_control_writer)
|
||||
|
||||
model_pred = self.denoising_unet(
|
||||
noisy_latents,
|
||||
timesteps,
|
||||
encoder_hidden_states=None,
|
||||
audio_cond_fea=audio_emb,
|
||||
).sample
|
||||
|
||||
return model_pred
|
||||
|
||||
|
||||
def get_noise_scheduler(cfg: argparse.Namespace):
|
||||
"""
|
||||
Create noise scheduler for training
|
||||
|
||||
Args:
|
||||
cfg (omegaconf.dictconfig.DictConfig): Configuration object.
|
||||
|
||||
Returns:
|
||||
train noise scheduler and val noise scheduler
|
||||
"""
|
||||
sched_kwargs = OmegaConf.to_container(cfg.noise_scheduler_kwargs)
|
||||
if cfg.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"})
|
||||
train_noise_scheduler = DDIMScheduler(**sched_kwargs)
|
||||
|
||||
return train_noise_scheduler, val_noise_scheduler
|
||||
|
||||
|
||||
def log_validation(
|
||||
vae,
|
||||
net,
|
||||
scheduler,
|
||||
accelerator,
|
||||
width,
|
||||
height,
|
||||
cfg,
|
||||
save_dir,
|
||||
global_step,
|
||||
val_dataset,
|
||||
weight_dtype
|
||||
):
|
||||
|
||||
ori_net = accelerator.unwrap_model(net)
|
||||
ori_net = copy.deepcopy(ori_net)
|
||||
reference_unet = ori_net.reference_unet
|
||||
denoising_unet = ori_net.denoising_unet
|
||||
|
||||
generator = torch.manual_seed(42)
|
||||
|
||||
pipe = StaticPipeline(
|
||||
vae=vae,
|
||||
reference_unet=reference_unet,
|
||||
denoising_unet=denoising_unet,
|
||||
scheduler=scheduler,
|
||||
)
|
||||
pipe = pipe.to("cuda")
|
||||
|
||||
pil_images = []
|
||||
|
||||
dataset_len = len(val_dataset)
|
||||
sample_idx = [np.random.randint(0, dataset_len-1) for _ in range(2)]
|
||||
for idx in sample_idx:
|
||||
sample = val_dataset[idx]
|
||||
|
||||
ref_image_pil = sample['ref_img']
|
||||
# gt_image_pil = Image.fromarray(sample['tar_gt']).convert("RGB")
|
||||
# ref_image = sample['ref_img'] # h*w*c
|
||||
# gt_images_no_trans = sample['tar_gt_no_trans'] # f*c*h*w
|
||||
gt_image = sample['img'] # f*c*h*w
|
||||
mask = sample['mask'].unsqueeze(0).to(device="cuda", dtype=weight_dtype) # b*c*f*h*w
|
||||
# mask_512 = sample['mask_512'].unsqueeze(0) # b*c*f*h*w
|
||||
audio_tensor = sample["audio_tensor"]
|
||||
|
||||
image = pipe(
|
||||
ref_image_pil,
|
||||
gt_image,
|
||||
audio_tensor,
|
||||
mask,
|
||||
width,
|
||||
height,
|
||||
20,
|
||||
3.5,
|
||||
generator=generator,
|
||||
).images # b*c*f*h*w
|
||||
|
||||
# gt_images_no_trans = gt_images_no_trans.permute(1,0,2,3).unsqueeze(0) # b*c*f*h*w
|
||||
# image = image * mask_512 + gt_images_no_trans * (1-mask_512)
|
||||
image = image[0, :, 0].permute(1, 2, 0).cpu().numpy() # (c, h, w)
|
||||
res_image_pil = Image.fromarray((image * 255).astype(np.uint8))
|
||||
# Save ref_image, src_image and the generated_image
|
||||
w, h = res_image_pil.size
|
||||
canvas = Image.new("RGB", (w * 3, h), "white")
|
||||
ref_image_pil = sample['ref_img'].convert("RGB").resize((w, h))
|
||||
gt_image_pil = sample['img'].convert("RGB").resize((w, h))
|
||||
|
||||
canvas.paste(ref_image_pil, (0, 0))
|
||||
canvas.paste(res_image_pil, (w, 0))
|
||||
canvas.paste(gt_image_pil, (w * 2, 0))
|
||||
|
||||
out_file = os.path.join(save_dir, f"{global_step:06d}-{idx}.jpg")
|
||||
canvas.save(out_file)
|
||||
|
||||
del pipe
|
||||
del ori_net
|
||||
torch.cuda.empty_cache()
|
||||
|
||||
return pil_images
|
||||
|
||||
|
||||
def train_stage1_process(cfg: argparse.Namespace) -> None:
|
||||
"""
|
||||
Trains the model using the given configuration (cfg).
|
||||
|
||||
Args:
|
||||
cfg (dict): The configuration dictionary containing the parameters for training.
|
||||
|
||||
Notes:
|
||||
- This function trains the model using the given configuration.
|
||||
- It initializes the necessary components for training, such as the pipeline, optimizer, and scheduler.
|
||||
- The training progress is logged and tracked using the accelerator.
|
||||
- The trained model is saved after the training is completed.
|
||||
"""
|
||||
kwargs = DistributedDataParallelKwargs(find_unused_parameters=True)
|
||||
accelerator = Accelerator(
|
||||
gradient_accumulation_steps=cfg.solver.gradient_accumulation_steps,
|
||||
mixed_precision=cfg.solver.mixed_precision,
|
||||
log_with="mlflow",
|
||||
project_dir="./mlruns",
|
||||
kwargs_handlers=[kwargs],
|
||||
)
|
||||
|
||||
# Make one log on every process with the configuration for debugging.
|
||||
logging.basicConfig(
|
||||
format="%(asctime)s - %(levelname)s - %(name)s - %(message)s",
|
||||
datefmt="%m/%d/%Y %H:%M:%S",
|
||||
level=logging.INFO,
|
||||
)
|
||||
|
||||
logger.info(accelerator.state, main_process_only=False)
|
||||
if accelerator.is_local_main_process:
|
||||
transformers.utils.logging.set_verbosity_warning()
|
||||
diffusers.utils.logging.set_verbosity_info()
|
||||
else:
|
||||
transformers.utils.logging.set_verbosity_error()
|
||||
diffusers.utils.logging.set_verbosity_error()
|
||||
|
||||
# If passed along, set the training seed now.
|
||||
if cfg.seed is not None:
|
||||
seed_everything(cfg.seed)
|
||||
|
||||
# create output dir for training
|
||||
exp_name = cfg.exp_name
|
||||
save_dir = f"{cfg.output_dir}/{exp_name}"
|
||||
checkpoint_dir = os.path.join(save_dir, "checkpoints")
|
||||
module_dir = os.path.join(save_dir, "modules")
|
||||
validation_dir = os.path.join(save_dir, "validation")
|
||||
|
||||
if accelerator.is_main_process:
|
||||
init_output_dir([save_dir, checkpoint_dir, module_dir, validation_dir])
|
||||
|
||||
accelerator.wait_for_everyone()
|
||||
|
||||
# create model
|
||||
if cfg.weight_dtype == "fp16":
|
||||
weight_dtype = torch.float16
|
||||
elif cfg.weight_dtype == "bf16":
|
||||
weight_dtype = torch.bfloat16
|
||||
elif cfg.weight_dtype == "fp32":
|
||||
weight_dtype = torch.float32
|
||||
else:
|
||||
raise ValueError(
|
||||
f"Do not support weight dtype: {cfg.weight_dtype} during training"
|
||||
)
|
||||
|
||||
# create model
|
||||
vae = AutoencoderKL.from_pretrained(cfg.vae_model_path).to(
|
||||
"cuda", dtype=weight_dtype
|
||||
)
|
||||
reference_unet = UNet2DConditionModel.from_pretrained(
|
||||
cfg.base_model_path,
|
||||
subfolder="unet",
|
||||
).to(device="cuda", dtype=weight_dtype)
|
||||
denoising_unet = EchoUNet3DConditionModel.from_pretrained_2d(
|
||||
cfg.base_model_path,
|
||||
"",
|
||||
subfolder="unet",
|
||||
unet_additional_kwargs={
|
||||
"use_motion_module": False,
|
||||
"unet_use_temporal_attention": False,
|
||||
"cross_attention_dim": cfg.cross_attention_dim
|
||||
}
|
||||
).to(device="cuda", dtype=weight_dtype)
|
||||
|
||||
denoising_unet.load_state_dict(
|
||||
torch.load(
|
||||
os.path.join(cfg.echo_path, "denoising_unet-220000.pth"),
|
||||
map_location="cpu",
|
||||
),
|
||||
strict=False,
|
||||
)
|
||||
reference_unet.load_state_dict(
|
||||
torch.load(
|
||||
os.path.join(cfg.echo_path, "reference_unet-220000.pth"),
|
||||
map_location="cpu",
|
||||
),
|
||||
strict=False,
|
||||
)
|
||||
|
||||
# Freeze
|
||||
vae.requires_grad_(False)
|
||||
denoising_unet.requires_grad_(True)
|
||||
reference_unet.requires_grad_(True)
|
||||
|
||||
reference_control_writer = ReferenceAttentionControl(
|
||||
reference_unet,
|
||||
do_classifier_free_guidance=False,
|
||||
mode="write",
|
||||
fusion_blocks="full",
|
||||
)
|
||||
reference_control_reader = ReferenceAttentionControl(
|
||||
denoising_unet,
|
||||
do_classifier_free_guidance=False,
|
||||
mode="read",
|
||||
fusion_blocks="full",
|
||||
audio_feature_ratio=3.0,
|
||||
)
|
||||
|
||||
net = Net(
|
||||
reference_unet,
|
||||
denoising_unet,
|
||||
reference_control_writer,
|
||||
reference_control_reader,
|
||||
).to(dtype=weight_dtype)
|
||||
|
||||
# get noise scheduler
|
||||
train_noise_scheduler, val_noise_scheduler = get_noise_scheduler(cfg)
|
||||
|
||||
# init optimizer
|
||||
if cfg.solver.enable_xformers_memory_efficient_attention:
|
||||
if is_xformers_available():
|
||||
reference_unet.enable_xformers_memory_efficient_attention()
|
||||
denoising_unet.enable_xformers_memory_efficient_attention()
|
||||
else:
|
||||
raise ValueError(
|
||||
"xformers is not available. Make sure it is installed correctly"
|
||||
)
|
||||
|
||||
if cfg.solver.gradient_checkpointing:
|
||||
reference_unet.enable_gradient_checkpointing()
|
||||
denoising_unet.enable_gradient_checkpointing()
|
||||
|
||||
if cfg.solver.scale_lr:
|
||||
learning_rate = (
|
||||
cfg.solver.learning_rate
|
||||
* cfg.solver.gradient_accumulation_steps
|
||||
* cfg.data.train_bs
|
||||
* accelerator.num_processes
|
||||
)
|
||||
else:
|
||||
learning_rate = cfg.solver.learning_rate
|
||||
|
||||
# Initialize the optimizer
|
||||
if cfg.solver.use_8bit_adam:
|
||||
try:
|
||||
import bitsandbytes as bnb
|
||||
except ImportError as exc:
|
||||
raise ImportError(
|
||||
"Please install bitsandbytes to use 8-bit Adam. You can do so by running `pip install bitsandbytes`"
|
||||
) from exc
|
||||
|
||||
optimizer_cls = bnb.optim.AdamW8bit
|
||||
else:
|
||||
optimizer_cls = torch.optim.AdamW
|
||||
|
||||
trainable_params = list(
|
||||
filter(lambda p: p.requires_grad, net.parameters()))
|
||||
optimizer = optimizer_cls(
|
||||
trainable_params,
|
||||
lr=learning_rate,
|
||||
betas=(cfg.solver.adam_beta1, cfg.solver.adam_beta2),
|
||||
weight_decay=cfg.solver.adam_weight_decay,
|
||||
eps=cfg.solver.adam_epsilon,
|
||||
)
|
||||
|
||||
# init scheduler
|
||||
lr_scheduler = get_scheduler(
|
||||
cfg.solver.lr_scheduler,
|
||||
optimizer=optimizer,
|
||||
num_warmup_steps=cfg.solver.lr_warmup_steps
|
||||
* cfg.solver.gradient_accumulation_steps,
|
||||
num_training_steps=cfg.solver.max_train_steps
|
||||
* cfg.solver.gradient_accumulation_steps,
|
||||
)
|
||||
|
||||
# get data loader
|
||||
train_dataset = FaceMaskDataset(
|
||||
img_size=(cfg.data.train_width, cfg.data.train_height),
|
||||
mask_path=cfg.data.mask_path,
|
||||
data_meta_path=cfg.data.train_path,
|
||||
sample_margin=cfg.data.sample_margin,
|
||||
)
|
||||
val_dataset = FaceMaskDatasetVal(
|
||||
img_size=(cfg.data.train_width, cfg.data.train_height),
|
||||
mask_path=cfg.data.mask_path,
|
||||
data_meta_path=cfg.data.val_path,
|
||||
sample_margin=cfg.data.sample_margin,
|
||||
)
|
||||
train_dataloader = torch.utils.data.DataLoader(
|
||||
train_dataset, batch_size=cfg.data.train_bs, shuffle=True, num_workers=4
|
||||
)
|
||||
|
||||
# Prepare everything with our `accelerator`.
|
||||
(
|
||||
net,
|
||||
optimizer,
|
||||
train_dataloader,
|
||||
lr_scheduler,
|
||||
) = accelerator.prepare(
|
||||
net,
|
||||
optimizer,
|
||||
train_dataloader,
|
||||
lr_scheduler,
|
||||
)
|
||||
# We need to recalculate our total training steps as the size of the training dataloader may have changed.
|
||||
num_update_steps_per_epoch = math.ceil(
|
||||
len(train_dataloader) / cfg.solver.gradient_accumulation_steps
|
||||
)
|
||||
# Afterwards we recalculate our number of training epochs
|
||||
num_train_epochs = math.ceil(
|
||||
cfg.solver.max_train_steps / num_update_steps_per_epoch
|
||||
)
|
||||
|
||||
# We need to initialize the trackers we use, and also store our configuration.
|
||||
# The trackers initializes automatically on the main process.
|
||||
if accelerator.is_main_process:
|
||||
run_time = datetime.now().strftime("%Y%m%d-%H%M")
|
||||
accelerator.init_trackers(
|
||||
cfg.exp_name,
|
||||
init_kwargs={"mlflow": {"run_name": run_time}},
|
||||
)
|
||||
# dump config file
|
||||
mlflow.log_dict(OmegaConf.to_container(cfg), "config.yaml")
|
||||
|
||||
logger.info(f"save config to {save_dir}")
|
||||
OmegaConf.save(
|
||||
cfg, os.path.join(save_dir, "config.yaml")
|
||||
)
|
||||
# Train!
|
||||
total_batch_size = (
|
||||
cfg.data.train_bs
|
||||
* accelerator.num_processes
|
||||
* cfg.solver.gradient_accumulation_steps
|
||||
)
|
||||
|
||||
logger.info("***** Running training *****")
|
||||
logger.info(f" Num examples = {len(train_dataset)}")
|
||||
logger.info(f" Num Epochs = {num_train_epochs}")
|
||||
logger.info(f" Instantaneous batch size per device = {cfg.data.train_bs}")
|
||||
logger.info(
|
||||
f" Total train batch size (w. parallel, distributed & accumulation) = {total_batch_size}"
|
||||
)
|
||||
logger.info(
|
||||
f" Gradient Accumulation steps = {cfg.solver.gradient_accumulation_steps}"
|
||||
)
|
||||
logger.info(f" Total optimization steps = {cfg.solver.max_train_steps}")
|
||||
global_step = 0
|
||||
first_epoch = 0
|
||||
|
||||
# load checkpoint
|
||||
# Potentially load in the weights and states from a previous save
|
||||
if cfg.resume_from_checkpoint:
|
||||
logger.info(f"Loading checkpoint from {checkpoint_dir}")
|
||||
global_step = load_checkpoint(cfg, checkpoint_dir, accelerator)
|
||||
first_epoch = global_step // num_update_steps_per_epoch
|
||||
|
||||
# Only show the progress bar once on each machine.
|
||||
progress_bar = tqdm(
|
||||
range(global_step, cfg.solver.max_train_steps),
|
||||
disable=not accelerator.is_main_process,
|
||||
)
|
||||
progress_bar.set_description("Steps")
|
||||
net.train()
|
||||
for _ in range(first_epoch, num_train_epochs):
|
||||
train_loss = 0.0
|
||||
for _, batch in enumerate(train_dataloader):
|
||||
with accelerator.accumulate(net):
|
||||
# Convert videos to latent space
|
||||
pixel_values = batch["img"].to(weight_dtype)
|
||||
with torch.no_grad():
|
||||
latents = vae.encode(pixel_values).latent_dist.sample()
|
||||
latents = latents.unsqueeze(2) # (b, c, 1, h, w)
|
||||
latents = latents * 0.18215
|
||||
|
||||
noise = torch.randn_like(latents)
|
||||
if cfg.noise_offset > 0.0:
|
||||
noise += cfg.noise_offset * torch.randn(
|
||||
(noise.shape[0], noise.shape[1], 1, 1, 1),
|
||||
device=noise.device,
|
||||
)
|
||||
|
||||
bsz = latents.shape[0]
|
||||
# Sample a random timestep for each video
|
||||
timesteps = torch.randint(
|
||||
0,
|
||||
train_noise_scheduler.num_train_timesteps,
|
||||
(bsz,),
|
||||
device=latents.device,
|
||||
)
|
||||
timesteps = timesteps.long()
|
||||
|
||||
uncond_fwd = np.random.random() < cfg.uncond_ratio
|
||||
|
||||
with torch.no_grad():
|
||||
ref_img = batch["ref_img"].to(
|
||||
dtype=vae.dtype, device=vae.device
|
||||
)
|
||||
ref_image_latents = vae.encode(
|
||||
ref_img
|
||||
).latent_dist.sample()
|
||||
ref_image_latents = ref_image_latents * 0.18215
|
||||
|
||||
mask = batch["mask"].to(device="cuda", dtype=weight_dtype)
|
||||
noise = noise * mask
|
||||
|
||||
# add noise
|
||||
noisy_latents = train_noise_scheduler.add_noise(
|
||||
latents, noise, timesteps
|
||||
)
|
||||
|
||||
# Get the target for loss depending on the prediction type
|
||||
if train_noise_scheduler.prediction_type == "epsilon":
|
||||
target = noise
|
||||
elif train_noise_scheduler.prediction_type == "v_prediction":
|
||||
target = train_noise_scheduler.get_velocity(
|
||||
latents, noise, timesteps
|
||||
)
|
||||
else:
|
||||
raise ValueError(
|
||||
f"Unknown prediction type {train_noise_scheduler.prediction_type}"
|
||||
)
|
||||
|
||||
model_pred = net(
|
||||
noisy_latents,
|
||||
timesteps,
|
||||
ref_image_latents,
|
||||
audio_emb=batch["audio_tensor"].to(
|
||||
dtype=weight_dtype),
|
||||
uncond_fwd=uncond_fwd,
|
||||
)
|
||||
|
||||
if cfg.snr_gamma == 0:
|
||||
loss = F.mse_loss(
|
||||
(model_pred*mask).float(), (target*mask).float(), reduction="mean"
|
||||
)
|
||||
else:
|
||||
snr = compute_snr(train_noise_scheduler, timesteps)
|
||||
if train_noise_scheduler.config.prediction_type == "v_prediction":
|
||||
# Velocity objective requires that we add one to SNR values before we divide by them.
|
||||
snr = snr + 1
|
||||
mse_loss_weights = (
|
||||
torch.stack(
|
||||
[snr, cfg.snr_gamma * torch.ones_like(timesteps)], dim=1
|
||||
).min(dim=1)[0]
|
||||
/ snr
|
||||
)
|
||||
loss = F.mse_loss(
|
||||
(model_pred*mask).float(), (target*mask).float(), reduction="none"
|
||||
)
|
||||
loss = (
|
||||
loss.mean(dim=list(range(1, len(loss.shape))))
|
||||
* mse_loss_weights
|
||||
)
|
||||
loss = loss.mean()
|
||||
|
||||
# Gather the losses across all processes for logging (if we use distributed training).
|
||||
avg_loss = accelerator.gather(
|
||||
loss.repeat(cfg.data.train_bs)).mean()
|
||||
train_loss += avg_loss.item() / cfg.solver.gradient_accumulation_steps
|
||||
|
||||
# Backpropagate
|
||||
accelerator.backward(loss*2)
|
||||
if accelerator.sync_gradients:
|
||||
accelerator.clip_grad_norm_(
|
||||
trainable_params,
|
||||
cfg.solver.max_grad_norm,
|
||||
)
|
||||
optimizer.step()
|
||||
lr_scheduler.step()
|
||||
optimizer.zero_grad()
|
||||
|
||||
if accelerator.sync_gradients:
|
||||
reference_control_reader.clear()
|
||||
reference_control_writer.clear()
|
||||
progress_bar.update(1)
|
||||
global_step += 1
|
||||
accelerator.log({"train_loss": train_loss}, step=global_step)
|
||||
train_loss = 0.0
|
||||
if global_step % cfg.checkpointing_steps == 0 or global_step == cfg.solver.max_train_steps:
|
||||
accelerator.wait_for_everyone()
|
||||
save_path = os.path.join(
|
||||
checkpoint_dir, f"checkpoint-{global_step}")
|
||||
if accelerator.is_main_process:
|
||||
delete_additional_ckpt(checkpoint_dir, 0)
|
||||
accelerator.save_state(save_path)
|
||||
accelerator.wait_for_everyone()
|
||||
unwrap_net = accelerator.unwrap_model(net)
|
||||
if accelerator.is_main_process:
|
||||
save_checkpoint(
|
||||
unwrap_net.reference_unet,
|
||||
module_dir,
|
||||
"reference_unet",
|
||||
global_step,
|
||||
total_limit=1,
|
||||
)
|
||||
save_checkpoint(
|
||||
unwrap_net.denoising_unet,
|
||||
module_dir,
|
||||
"denoising_unet",
|
||||
global_step,
|
||||
total_limit=1,
|
||||
)
|
||||
|
||||
if global_step % cfg.val.validation_steps == 0:
|
||||
if accelerator.is_main_process:
|
||||
generator = torch.Generator(device=accelerator.device)
|
||||
generator.manual_seed(cfg.seed)
|
||||
log_validation(
|
||||
vae=vae,
|
||||
net=net,
|
||||
scheduler=val_noise_scheduler,
|
||||
accelerator=accelerator,
|
||||
width=cfg.data.train_width,
|
||||
height=cfg.data.train_height,
|
||||
cfg=cfg,
|
||||
save_dir=validation_dir,
|
||||
global_step=global_step,
|
||||
val_dataset=val_dataset,
|
||||
weight_dtype=weight_dtype
|
||||
)
|
||||
|
||||
logs = {
|
||||
"step_loss": loss.detach().item(),
|
||||
"lr": lr_scheduler.get_last_lr()[0],
|
||||
}
|
||||
progress_bar.set_postfix(**logs)
|
||||
|
||||
writer.add_scalar("loss", loss.detach().item(), global_step)
|
||||
|
||||
if global_step >= cfg.solver.max_train_steps:
|
||||
# process final module weight for stage2
|
||||
if accelerator.is_main_process:
|
||||
move_final_checkpoint(save_dir, module_dir, "reference_unet")
|
||||
move_final_checkpoint(save_dir, module_dir, "denoising_unet")
|
||||
break
|
||||
|
||||
accelerator.wait_for_everyone()
|
||||
accelerator.end_training()
|
||||
|
||||
|
||||
def load_config(config_path: str) -> dict:
|
||||
"""
|
||||
Loads the configuration file.
|
||||
|
||||
Args:
|
||||
config_path (str): Path to the configuration file.
|
||||
|
||||
Returns:
|
||||
dict: The configuration dictionary.
|
||||
"""
|
||||
|
||||
if config_path.endswith(".yaml"):
|
||||
return OmegaConf.load(config_path)
|
||||
if config_path.endswith(".py"):
|
||||
return import_filename(config_path).cfg
|
||||
raise ValueError("Unsupported format for config file")
|
||||
|
||||
|
||||
if __name__ == "__main__":
|
||||
parser = argparse.ArgumentParser()
|
||||
parser.add_argument("--config", type=str,
|
||||
default="./configs/train/stage1.yaml")
|
||||
args = parser.parse_args()
|
||||
|
||||
config = load_config(args.config)
|
||||
train_stage1_process(config)
|
||||
@@ -0,0 +1,855 @@
|
||||
# pylint: disable=E1101,C0415,W0718,R0801
|
||||
# scripts/train_stage2.py
|
||||
"""
|
||||
This is the main training script for stage 2 of the project.
|
||||
It imports necessary packages, defines necessary classes and functions, and trains the model using the provided configuration.
|
||||
|
||||
The script includes the following classes and functions:
|
||||
|
||||
1. Net: A PyTorch model that takes noisy latents, timesteps, reference image latents, face embeddings,
|
||||
and face masks as input and returns the denoised latents.
|
||||
2. get_attention_mask: A function that rearranges the mask tensors to the required format.
|
||||
3. get_noise_scheduler: A function that creates and returns the noise schedulers for training and validation.
|
||||
4. process_audio_emb: A function that processes the audio embeddings to concatenate with other tensors.
|
||||
5. log_validation: A function that logs the validation information using the given VAE, image encoder,
|
||||
network, scheduler, accelerator, width, height, and configuration.
|
||||
6. train_stage2_process: A function that processes the training stage 2 using the given configuration.
|
||||
7. load_config: A function that loads the configuration file from the given path.
|
||||
|
||||
The script also includes the necessary imports and a brief description of the purpose of the file.
|
||||
"""
|
||||
|
||||
import argparse
|
||||
import copy
|
||||
import logging
|
||||
import math
|
||||
import os
|
||||
import random
|
||||
import time
|
||||
import warnings
|
||||
from datetime import datetime
|
||||
from typing import List, Tuple
|
||||
|
||||
import diffusers
|
||||
import mlflow
|
||||
import torch
|
||||
import torch.nn.functional as F
|
||||
import torch.utils.checkpoint
|
||||
import transformers
|
||||
from accelerate import Accelerator
|
||||
from accelerate.logging import get_logger
|
||||
from accelerate.utils import DistributedDataParallelKwargs
|
||||
from diffusers import AutoencoderKL, DDIMScheduler
|
||||
from diffusers.optimization import get_scheduler
|
||||
from diffusers.utils import check_min_version
|
||||
from diffusers.utils.import_utils import is_xformers_available
|
||||
from einops import rearrange, repeat
|
||||
from omegaconf import OmegaConf
|
||||
import numpy as np
|
||||
from torch import nn
|
||||
from tqdm.auto import tqdm
|
||||
from torch.utils.tensorboard import SummaryWriter
|
||||
import torch.nn.init as init
|
||||
|
||||
from .src.pipelines.face_animate import FaceAnimatePipeline
|
||||
from .src.datasets.talk_video import TalkingVideoDataset, TalkingVideoDatasetVal
|
||||
from .src.models.mutual_self_attention import ReferenceAttentionControl
|
||||
from .src.models.unet_2d_condition import UNet2DConditionModel
|
||||
from .src.models.unet_3d_echo import EchoUNet3DConditionModel
|
||||
from .src.utils.util import (compute_snr, delete_additional_ckpt,
|
||||
import_filename, init_output_dir,
|
||||
load_checkpoint, save_checkpoint,
|
||||
seed_everything, tensor_to_video_new)
|
||||
|
||||
warnings.filterwarnings("ignore")
|
||||
|
||||
# Will error if the minimal version of diffusers is not installed. Remove at your own risks.
|
||||
check_min_version("0.10.0.dev0")
|
||||
|
||||
logger = get_logger(__name__, log_level="INFO")
|
||||
|
||||
writer = SummaryWriter()
|
||||
|
||||
class Net(nn.Module):
|
||||
"""
|
||||
The Net class defines a neural network model that combines a reference UNet2DConditionModel,
|
||||
a denoising UNet3DConditionModel, a face locator, and other components to animate a face in a static image.
|
||||
|
||||
Args:
|
||||
reference_unet (UNet2DConditionModel): The reference UNet2DConditionModel used for face animation.
|
||||
denoising_unet (UNet3DConditionModel): The denoising UNet3DConditionModel used for face animation.
|
||||
face_locator (FaceLocator): The face locator model used for face animation.
|
||||
reference_control_writer: The reference control writer component.
|
||||
reference_control_reader: The reference control reader component.
|
||||
imageproj: The image projection model.
|
||||
audioproj: The audio projection model.
|
||||
|
||||
Forward method:
|
||||
noisy_latents (torch.Tensor): The noisy latents tensor.
|
||||
timesteps (torch.Tensor): The timesteps tensor.
|
||||
ref_image_latents (torch.Tensor): The reference image latents tensor.
|
||||
face_emb (torch.Tensor): The face embeddings tensor.
|
||||
audio_emb (torch.Tensor): The audio embeddings tensor.
|
||||
mask (torch.Tensor): Hard face mask for face locator.
|
||||
full_mask (torch.Tensor): Pose Mask.
|
||||
face_mask (torch.Tensor): Face Mask
|
||||
lip_mask (torch.Tensor): Lip Mask
|
||||
uncond_img_fwd (bool): A flag indicating whether to perform reference image unconditional forward pass.
|
||||
uncond_audio_fwd (bool): A flag indicating whether to perform audio unconditional forward pass.
|
||||
|
||||
Returns:
|
||||
torch.Tensor: The output tensor of the neural network model.
|
||||
"""
|
||||
def __init__(
|
||||
self,
|
||||
reference_unet: UNet2DConditionModel,
|
||||
denoising_unet: EchoUNet3DConditionModel,
|
||||
reference_control_writer,
|
||||
reference_control_reader,
|
||||
):
|
||||
super().__init__()
|
||||
self.reference_unet = reference_unet
|
||||
self.denoising_unet = denoising_unet
|
||||
self.reference_control_writer = reference_control_writer
|
||||
self.reference_control_reader = reference_control_reader
|
||||
|
||||
def forward(
|
||||
self,
|
||||
noisy_latents: torch.Tensor,
|
||||
timesteps: torch.Tensor,
|
||||
ref_image_latents: torch.Tensor,
|
||||
audio_emb: torch.Tensor,
|
||||
uncond_img_fwd: bool = False,
|
||||
uncond_audio_fwd: bool = False,
|
||||
):
|
||||
"""
|
||||
simple docstring to prevent pylint error
|
||||
"""
|
||||
audio_emb = audio_emb.to(
|
||||
device=self.reference_unet.device, dtype=self.reference_unet.dtype)
|
||||
|
||||
# condition forward
|
||||
if not uncond_img_fwd:
|
||||
ref_timesteps = torch.zeros_like(timesteps)
|
||||
ref_timesteps = repeat(
|
||||
ref_timesteps,
|
||||
"b -> (repeat b)",
|
||||
repeat=ref_image_latents.size(0) // ref_timesteps.size(0),
|
||||
)
|
||||
self.reference_unet(
|
||||
ref_image_latents,
|
||||
ref_timesteps,
|
||||
encoder_hidden_states=None,
|
||||
return_dict=False,
|
||||
)
|
||||
self.reference_control_reader.update(self.reference_control_writer)
|
||||
|
||||
if uncond_audio_fwd:
|
||||
audio_emb = torch.zeros_like(audio_emb).to(
|
||||
device=audio_emb.device, dtype=audio_emb.dtype
|
||||
)
|
||||
|
||||
model_pred = self.denoising_unet(
|
||||
noisy_latents,
|
||||
timesteps,
|
||||
encoder_hidden_states=None,
|
||||
audio_cond_fea=audio_emb,
|
||||
).sample
|
||||
|
||||
return model_pred
|
||||
|
||||
|
||||
def get_noise_scheduler(cfg: argparse.Namespace) -> Tuple[DDIMScheduler, DDIMScheduler]:
|
||||
"""
|
||||
Create noise scheduler for training.
|
||||
|
||||
Args:
|
||||
cfg (argparse.Namespace): Configuration object.
|
||||
|
||||
Returns:
|
||||
Tuple[DDIMScheduler, DDIMScheduler]: Train noise scheduler and validation noise scheduler.
|
||||
"""
|
||||
|
||||
sched_kwargs = OmegaConf.to_container(cfg.noise_scheduler_kwargs)
|
||||
if cfg.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"})
|
||||
train_noise_scheduler = DDIMScheduler(**sched_kwargs)
|
||||
|
||||
return train_noise_scheduler, val_noise_scheduler
|
||||
|
||||
|
||||
def log_validation(
|
||||
accelerator: Accelerator,
|
||||
vae: AutoencoderKL,
|
||||
net: Net,
|
||||
scheduler: DDIMScheduler,
|
||||
width: int,
|
||||
height: int,
|
||||
clip_length: int = 14,
|
||||
generator: torch.Generator = None,
|
||||
cfg: dict = None,
|
||||
save_dir: str = None,
|
||||
global_step: int = 0,
|
||||
times: int = None,
|
||||
val_dataset = None,
|
||||
weight_dtype = None
|
||||
) -> None:
|
||||
"""
|
||||
Log validation video during the training process.
|
||||
|
||||
Args:
|
||||
accelerator (Accelerator): The accelerator for distributed training.
|
||||
vae (AutoencoderKL): The autoencoder model.
|
||||
net (Net): The main neural network model.
|
||||
scheduler (DDIMScheduler): The scheduler for noise.
|
||||
width (int): The width of the input images.
|
||||
height (int): The height of the input images.
|
||||
clip_length (int): The length of the video clips. Defaults to 24.
|
||||
generator (torch.Generator): The random number generator. Defaults to None.
|
||||
cfg (dict): The configuration dictionary. Defaults to None.
|
||||
save_dir (str): The directory to save validation results. Defaults to None.
|
||||
global_step (int): The current global step in training. Defaults to 0.
|
||||
times (int): The number of inference times. Defaults to None.
|
||||
face_analysis_model_path (str): The path to the face analysis model. Defaults to "".
|
||||
|
||||
Returns:
|
||||
torch.Tensor: The tensor result of the validation.
|
||||
"""
|
||||
|
||||
generator = torch.manual_seed(42)
|
||||
|
||||
pipeline = FaceAnimatePipeline(
|
||||
vae=vae,
|
||||
reference_unet=accelerator.unwrap_model(net).reference_unet,
|
||||
denoising_unet=accelerator.unwrap_model(net).denoising_unet,
|
||||
scheduler=scheduler,
|
||||
)
|
||||
pipeline = pipeline.to("cuda")
|
||||
|
||||
dataset_len = len(val_dataset)
|
||||
sample_idx = [np.random.randint(0, dataset_len-1) for _ in range(1)]
|
||||
|
||||
for idx in sample_idx:
|
||||
sample = val_dataset[idx]
|
||||
same = sample["same"]
|
||||
video_path = sample["video_dir"]
|
||||
audio_path = sample["audio_path"]
|
||||
audio_emb = sample["audio_tensor"]
|
||||
|
||||
target_images = sample["pixel_values_vid"].to(device="cuda", dtype=weight_dtype)
|
||||
mask = sample['mask'].unsqueeze(0).to(device="cuda", dtype=weight_dtype) # b*c*f*h*w
|
||||
|
||||
min_time = min(audio_emb.shape[0], target_images.shape[0], 5*25)
|
||||
target_images = target_images[:min_time]
|
||||
audio_emb = audio_emb[:min_time]
|
||||
tensor_result = []
|
||||
generator = torch.manual_seed(42)
|
||||
|
||||
audio_emb = audio_emb.unsqueeze(0).to(device="cuda", dtype=weight_dtype)
|
||||
|
||||
try:
|
||||
pipeline_output = pipeline(
|
||||
gt_encode=target_images,
|
||||
audio_tensor=audio_emb,
|
||||
mask_main=mask,
|
||||
width=cfg.data.train_width,
|
||||
height=cfg.data.train_height,
|
||||
video_length=min_time,
|
||||
num_inference_steps=cfg.inference_steps,
|
||||
guidance_scale=cfg.cfg_scale,
|
||||
generator=generator,
|
||||
)
|
||||
except Exception as e:
|
||||
import traceback
|
||||
print("Error encountered during training step:")
|
||||
traceback.print_exc() # Prints the full error traceback
|
||||
exit(1)
|
||||
|
||||
tensor_result.append(pipeline_output.videos)
|
||||
|
||||
tensor_result = torch.cat(tensor_result, dim=2)
|
||||
tensor_result = tensor_result.squeeze(0)
|
||||
tensor_result = tensor_result[:, :]
|
||||
video_name = os.path.basename(video_path).split('.')[-2]
|
||||
output_file = os.path.join(save_dir,f"{global_step}_{video_name}_{idx}_{same}.mp4")
|
||||
# save the result after all iteration
|
||||
tensor_to_video_new(tensor_result, output_file, audio_path)
|
||||
|
||||
|
||||
# clean up
|
||||
# del tmp_denoising_unet
|
||||
del pipeline
|
||||
torch.cuda.empty_cache()
|
||||
|
||||
return tensor_result
|
||||
|
||||
|
||||
def train_stage2_process(cfg: argparse.Namespace) -> None:
|
||||
"""
|
||||
Trains the model using the given configuration (cfg).
|
||||
|
||||
Args:
|
||||
cfg (dict): The configuration dictionary containing the parameters for training.
|
||||
|
||||
Notes:
|
||||
- This function trains the model using the given configuration.
|
||||
- It initializes the necessary components for training, such as the pipeline, optimizer, and scheduler.
|
||||
- The training progress is logged and tracked using the accelerator.
|
||||
- The trained model is saved after the training is completed.
|
||||
"""
|
||||
kwargs = DistributedDataParallelKwargs(find_unused_parameters=False)
|
||||
accelerator = Accelerator(
|
||||
gradient_accumulation_steps=cfg.solver.gradient_accumulation_steps,
|
||||
mixed_precision=cfg.solver.mixed_precision,
|
||||
log_with="mlflow",
|
||||
project_dir="./mlruns",
|
||||
kwargs_handlers=[kwargs],
|
||||
)
|
||||
|
||||
# Make one log on every process with the configuration for debugging.
|
||||
logging.basicConfig(
|
||||
format="%(asctime)s - %(levelname)s - %(name)s - %(message)s",
|
||||
datefmt="%m/%d/%Y %H:%M:%S",
|
||||
level=logging.INFO,
|
||||
)
|
||||
logger.info(accelerator.state, main_process_only=False)
|
||||
if accelerator.is_local_main_process:
|
||||
transformers.utils.logging.set_verbosity_warning()
|
||||
diffusers.utils.logging.set_verbosity_info()
|
||||
else:
|
||||
transformers.utils.logging.set_verbosity_error()
|
||||
diffusers.utils.logging.set_verbosity_error()
|
||||
|
||||
# If passed along, set the training seed now.
|
||||
if cfg.seed is not None:
|
||||
seed_everything(cfg.seed)
|
||||
|
||||
# create output dir for training
|
||||
exp_name = cfg.exp_name
|
||||
save_dir = f"{cfg.output_dir}/{exp_name}"
|
||||
checkpoint_dir = os.path.join(save_dir, "checkpoints")
|
||||
module_dir = os.path.join(save_dir, "modules")
|
||||
validation_dir = os.path.join(save_dir, "validation")
|
||||
if accelerator.is_main_process:
|
||||
init_output_dir([save_dir, checkpoint_dir, module_dir, validation_dir])
|
||||
|
||||
accelerator.wait_for_everyone()
|
||||
|
||||
if cfg.weight_dtype == "fp16":
|
||||
weight_dtype = torch.float16
|
||||
elif cfg.weight_dtype == "bf16":
|
||||
weight_dtype = torch.bfloat16
|
||||
elif cfg.weight_dtype == "fp32":
|
||||
weight_dtype = torch.float32
|
||||
else:
|
||||
raise ValueError(
|
||||
f"Do not support weight dtype: {cfg.weight_dtype} during training"
|
||||
)
|
||||
|
||||
# Create Models
|
||||
vae = AutoencoderKL.from_pretrained(cfg.vae_model_path).to(
|
||||
"cuda", dtype=weight_dtype
|
||||
)
|
||||
reference_unet = UNet2DConditionModel.from_pretrained(
|
||||
cfg.base_model_path,
|
||||
subfolder="unet",
|
||||
).to(device="cuda", dtype=weight_dtype)
|
||||
denoising_unet = EchoUNet3DConditionModel.from_pretrained_2d(
|
||||
cfg.base_model_path,
|
||||
cfg.mm_path,
|
||||
subfolder="unet",
|
||||
unet_additional_kwargs=OmegaConf.to_container(
|
||||
cfg.unet_additional_kwargs)
|
||||
).to(device="cuda", dtype=weight_dtype)
|
||||
|
||||
# load module weight from stage 1
|
||||
stage1_ckpt_dir = cfg.stage1_ckpt_dir
|
||||
denoising_unet.load_state_dict(
|
||||
torch.load(
|
||||
os.path.join(stage1_ckpt_dir, "denoising_unet-44000.pth"),
|
||||
map_location="cpu",
|
||||
),
|
||||
strict=False,
|
||||
)
|
||||
reference_unet.load_state_dict(
|
||||
torch.load(
|
||||
os.path.join(stage1_ckpt_dir, "reference_unet-44000.pth"),
|
||||
map_location="cpu",
|
||||
),
|
||||
strict=False,
|
||||
)
|
||||
|
||||
# Freeze
|
||||
vae.requires_grad_(False)
|
||||
reference_unet.requires_grad_(False)
|
||||
denoising_unet.requires_grad_(False)
|
||||
|
||||
# Set motion module learnable
|
||||
trainable_modules = cfg.trainable_para
|
||||
for name, module in denoising_unet.named_modules():
|
||||
if any(trainable_mod in name for trainable_mod in trainable_modules):
|
||||
for param in module.parameters():
|
||||
param.requires_grad_(True)
|
||||
|
||||
# audio attention parames of denoising_unet should be updated
|
||||
for name, param in denoising_unet.named_parameters():
|
||||
if ("attn2" in name) or ("norm2" in name):
|
||||
param.requires_grad_(True)
|
||||
|
||||
reference_control_writer = ReferenceAttentionControl(
|
||||
reference_unet,
|
||||
do_classifier_free_guidance=False,
|
||||
mode="write",
|
||||
fusion_blocks="full",
|
||||
)
|
||||
reference_control_reader = ReferenceAttentionControl(
|
||||
denoising_unet,
|
||||
do_classifier_free_guidance=False,
|
||||
mode="read",
|
||||
fusion_blocks="full",
|
||||
audio_feature_ratio=3.0,
|
||||
)
|
||||
|
||||
net = Net(
|
||||
reference_unet,
|
||||
denoising_unet,
|
||||
reference_control_writer,
|
||||
reference_control_reader,
|
||||
).to(dtype=weight_dtype)
|
||||
|
||||
# unwrap_net = accelerator.unwrap_model(net)
|
||||
# print(f"mm denoise : {net.denoising_unet.state_dict()['down_blocks.0.motion_modules.0.temporal_transformer.transformer_blocks.0.attention_blocks.1.to_v.weight']}")
|
||||
|
||||
if cfg.resume:
|
||||
denoising_unet.load_state_dict(
|
||||
torch.load(
|
||||
cfg.resume_denoise_path,
|
||||
map_location="cpu",
|
||||
),
|
||||
strict=False,
|
||||
)
|
||||
reference_unet.load_state_dict(
|
||||
torch.load(
|
||||
cfg.resume_reference_path,
|
||||
map_location="cpu",
|
||||
),
|
||||
strict=False,
|
||||
)
|
||||
print("loaded weight from resume model")
|
||||
else:
|
||||
#load weight of net for just audio and motion
|
||||
pretrained_model = torch.load(os.path.join(cfg.echo_path, "denoising_unet.pth"),map_location="cpu",)
|
||||
weight_modules = ["motion_modules"]
|
||||
|
||||
filtered_state_dict = {}
|
||||
for k, v in pretrained_model.items():
|
||||
if any(trainable_mod in k for trainable_mod in weight_modules):
|
||||
filtered_state_dict[k] = v
|
||||
|
||||
m , u = net.denoising_unet.load_state_dict(filtered_state_dict, strict=False)
|
||||
|
||||
|
||||
# assert len(m) == 0 and len(u) == 0, "Fail to load correct checkpoint."
|
||||
print("loaded weight from Echo")
|
||||
|
||||
# get noise scheduler
|
||||
train_noise_scheduler, val_noise_scheduler = get_noise_scheduler(cfg)
|
||||
|
||||
if cfg.solver.enable_xformers_memory_efficient_attention:
|
||||
if is_xformers_available():
|
||||
reference_unet.enable_xformers_memory_efficient_attention()
|
||||
denoising_unet.enable_xformers_memory_efficient_attention()
|
||||
|
||||
else:
|
||||
raise ValueError(
|
||||
"xformers is not available. Make sure it is installed correctly"
|
||||
)
|
||||
|
||||
if cfg.solver.gradient_checkpointing:
|
||||
reference_unet.enable_gradient_checkpointing()
|
||||
denoising_unet.enable_gradient_checkpointing()
|
||||
|
||||
if cfg.solver.scale_lr:
|
||||
learning_rate = (
|
||||
cfg.solver.learning_rate
|
||||
* cfg.solver.gradient_accumulation_steps
|
||||
* cfg.data.train_bs
|
||||
* accelerator.num_processes
|
||||
)
|
||||
else:
|
||||
learning_rate = cfg.solver.learning_rate
|
||||
|
||||
# Initialize the optimizer
|
||||
if cfg.solver.use_8bit_adam:
|
||||
try:
|
||||
import bitsandbytes as bnb
|
||||
except ImportError as exc:
|
||||
raise ImportError(
|
||||
"Please install bitsandbytes to use 8-bit Adam. You can do so by running `pip install bitsandbytes`"
|
||||
) from exc
|
||||
optimizer_cls = bnb.optim.AdamW8bit
|
||||
else:
|
||||
optimizer_cls = torch.optim.AdamW
|
||||
|
||||
trainable_params = list(
|
||||
filter(lambda p: p.requires_grad, net.parameters()))
|
||||
logger.info(f"Total trainable params {len(trainable_params)}")
|
||||
optimizer = optimizer_cls(
|
||||
trainable_params,
|
||||
lr=learning_rate,
|
||||
betas=(cfg.solver.adam_beta1, cfg.solver.adam_beta2),
|
||||
weight_decay=cfg.solver.adam_weight_decay,
|
||||
eps=cfg.solver.adam_epsilon,
|
||||
)
|
||||
|
||||
# Scheduler
|
||||
lr_scheduler = get_scheduler(
|
||||
cfg.solver.lr_scheduler,
|
||||
optimizer=optimizer,
|
||||
num_warmup_steps=cfg.solver.lr_warmup_steps
|
||||
* cfg.solver.gradient_accumulation_steps,
|
||||
num_training_steps=cfg.solver.max_train_steps
|
||||
* cfg.solver.gradient_accumulation_steps,
|
||||
)
|
||||
|
||||
# get data loader
|
||||
train_dataset = TalkingVideoDataset(
|
||||
img_size=(cfg.data.train_width, cfg.data.train_height),
|
||||
mask_path=cfg.data.mask_path,
|
||||
sample_rate=cfg.data.sample_rate,
|
||||
n_sample_frames=cfg.data.n_sample_frames,
|
||||
audio_margin=cfg.data.audio_margin,
|
||||
data_meta_path=cfg.data.train_path,
|
||||
wav2vec_cfg=cfg.wav2vec_config,
|
||||
)
|
||||
|
||||
val_dataset = TalkingVideoDatasetVal(
|
||||
img_size=(cfg.data.train_width, cfg.data.train_height),
|
||||
mask_path=cfg.data.mask_path,
|
||||
sample_rate=cfg.data.sample_rate,
|
||||
n_sample_frames=cfg.data.n_sample_frames,
|
||||
audio_margin=cfg.data.audio_margin,
|
||||
data_meta_path=cfg.data.val_path,
|
||||
wav2vec_cfg=cfg.wav2vec_config,
|
||||
)
|
||||
train_dataloader = torch.utils.data.DataLoader(
|
||||
train_dataset, batch_size=cfg.data.train_bs, shuffle=True, num_workers=16
|
||||
)
|
||||
|
||||
# Prepare everything with our `accelerator`.
|
||||
(
|
||||
net,
|
||||
optimizer,
|
||||
train_dataloader,
|
||||
lr_scheduler,
|
||||
) = accelerator.prepare(
|
||||
net,
|
||||
optimizer,
|
||||
train_dataloader,
|
||||
lr_scheduler,
|
||||
)
|
||||
|
||||
# We need to recalculate our total training steps as the size of the training dataloader may have changed.
|
||||
num_update_steps_per_epoch = math.ceil(
|
||||
len(train_dataloader) / cfg.solver.gradient_accumulation_steps
|
||||
)
|
||||
# Afterwards we recalculate our number of training epochs
|
||||
num_train_epochs = math.ceil(
|
||||
cfg.solver.max_train_steps / num_update_steps_per_epoch
|
||||
)
|
||||
|
||||
# We need to initialize the trackers we use, and also store our configuration.
|
||||
# The trackers initializes automatically on the main process.
|
||||
if accelerator.is_main_process:
|
||||
run_time = datetime.now().strftime("%Y%m%d-%H%M")
|
||||
accelerator.init_trackers(
|
||||
exp_name,
|
||||
init_kwargs={"mlflow": {"run_name": run_time}},
|
||||
)
|
||||
# dump config file
|
||||
mlflow.log_dict(
|
||||
OmegaConf.to_container(
|
||||
cfg), "config.yaml"
|
||||
)
|
||||
logger.info(f"save config to {save_dir}")
|
||||
OmegaConf.save(
|
||||
cfg, os.path.join(save_dir, "config.yaml")
|
||||
)
|
||||
|
||||
# Train!
|
||||
total_batch_size = (
|
||||
cfg.data.train_bs
|
||||
* accelerator.num_processes
|
||||
* cfg.solver.gradient_accumulation_steps
|
||||
)
|
||||
|
||||
logger.info("***** Running training *****")
|
||||
logger.info(f" Num examples = {len(train_dataset)}")
|
||||
logger.info(f" Num Epochs = {num_train_epochs}")
|
||||
logger.info(f" Instantaneous batch size per device = {cfg.data.train_bs}")
|
||||
logger.info(
|
||||
f" Total train batch size (w. parallel, distributed & accumulation) = {total_batch_size}"
|
||||
)
|
||||
logger.info(
|
||||
f" Gradient Accumulation steps = {cfg.solver.gradient_accumulation_steps}"
|
||||
)
|
||||
logger.info(f" Total optimization steps = {cfg.solver.max_train_steps}")
|
||||
global_step = 8000
|
||||
first_epoch = 0
|
||||
|
||||
# # Potentially load in the weights and states from a previous save
|
||||
if cfg.resume_from_checkpoint:
|
||||
logger.info(f"Loading checkpoint from {checkpoint_dir}")
|
||||
global_step = load_checkpoint(cfg, checkpoint_dir, accelerator)
|
||||
first_epoch = global_step // num_update_steps_per_epoch
|
||||
|
||||
# Only show the progress bar once on each machine.
|
||||
progress_bar = tqdm(
|
||||
range(global_step, cfg.solver.max_train_steps),
|
||||
disable=not accelerator.is_local_main_process,
|
||||
)
|
||||
progress_bar.set_description("Steps")
|
||||
|
||||
|
||||
for _ in range(first_epoch, num_train_epochs):
|
||||
train_loss = 0.0
|
||||
for _, batch in enumerate(train_dataloader):
|
||||
with accelerator.accumulate(net):
|
||||
# Convert videos to latent space
|
||||
pixel_values_vid = batch["pixel_values_vid"].to(weight_dtype)
|
||||
|
||||
with torch.no_grad():
|
||||
video_length = pixel_values_vid.shape[1]
|
||||
pixel_values_vid = rearrange(
|
||||
pixel_values_vid, "b f c h w -> (b f) c h w"
|
||||
)
|
||||
latents = vae.encode(pixel_values_vid).latent_dist.sample()
|
||||
latents = rearrange(
|
||||
latents, "(b f) c h w -> b c f h w", f=video_length
|
||||
)
|
||||
latents = latents * 0.18215
|
||||
|
||||
noise = torch.randn_like(latents)
|
||||
if cfg.noise_offset > 0:
|
||||
noise += cfg.noise_offset * torch.randn(
|
||||
(latents.shape[0], latents.shape[1], 1, 1, 1),
|
||||
device=latents.device,
|
||||
)
|
||||
|
||||
bsz = latents.shape[0]
|
||||
# Sample a random timestep for each video
|
||||
timesteps = torch.randint(
|
||||
0,
|
||||
train_noise_scheduler.num_train_timesteps,
|
||||
(bsz,),
|
||||
device=latents.device,
|
||||
)
|
||||
timesteps = timesteps.long()
|
||||
|
||||
uncond_img_fwd = np.random.random() < cfg.uncond_img_ratio
|
||||
uncond_audio_fwd = np.random.random() < cfg.uncond_audio_ratio
|
||||
|
||||
pixel_values_ref_img = batch["pixel_values_ref_img"].to(
|
||||
dtype=weight_dtype
|
||||
)
|
||||
|
||||
ref_img_and_motion = rearrange(
|
||||
pixel_values_ref_img, "b f c h w -> (b f) c h w"
|
||||
)
|
||||
|
||||
with torch.no_grad():
|
||||
ref_image_latents = vae.encode(
|
||||
ref_img_and_motion
|
||||
).latent_dist.sample()
|
||||
ref_image_latents = ref_image_latents * 0.18215
|
||||
|
||||
mask = batch["mask"].to(device="cuda", dtype=weight_dtype)
|
||||
noise = noise * mask
|
||||
|
||||
uncond_target_fwd = np.random.random() < cfg.uncond_target_ratio
|
||||
if uncond_target_fwd:
|
||||
latents = latents * mask
|
||||
|
||||
# add noise
|
||||
noisy_latents = train_noise_scheduler.add_noise(
|
||||
latents, noise, timesteps
|
||||
)
|
||||
|
||||
# Get the target for loss depending on the prediction type
|
||||
if train_noise_scheduler.prediction_type == "epsilon":
|
||||
target = noise
|
||||
elif train_noise_scheduler.prediction_type == "v_prediction":
|
||||
target = train_noise_scheduler.get_velocity(
|
||||
latents, noise, timesteps
|
||||
)
|
||||
else:
|
||||
raise ValueError(
|
||||
f"Unknown prediction type {train_noise_scheduler.prediction_type}"
|
||||
)
|
||||
|
||||
# ---- Forward!!! -----
|
||||
model_pred = net(
|
||||
noisy_latents=noisy_latents,
|
||||
timesteps=timesteps,
|
||||
ref_image_latents=ref_image_latents,
|
||||
audio_emb=batch["audio_tensor"].to(
|
||||
dtype=weight_dtype),
|
||||
uncond_img_fwd=uncond_img_fwd,
|
||||
uncond_audio_fwd=uncond_audio_fwd,
|
||||
)
|
||||
|
||||
if cfg.snr_gamma == 0:
|
||||
loss = F.mse_loss(
|
||||
(model_pred*mask).float(),
|
||||
(target*mask).float(),
|
||||
reduction="mean",
|
||||
)
|
||||
else:
|
||||
snr = compute_snr(train_noise_scheduler, timesteps)
|
||||
if train_noise_scheduler.config.prediction_type == "v_prediction":
|
||||
# Velocity objective requires that we add one to SNR values before we divide by them.
|
||||
snr = snr + 1
|
||||
mse_loss_weights = (
|
||||
torch.stack(
|
||||
[snr, cfg.snr_gamma * torch.ones_like(timesteps)], dim=1
|
||||
).min(dim=1)[0]
|
||||
/ snr
|
||||
)
|
||||
loss = F.mse_loss(
|
||||
(model_pred*mask).float(),
|
||||
(target*mask).float(),
|
||||
reduction="mean",
|
||||
)
|
||||
loss = (
|
||||
loss.mean(dim=list(range(1, len(loss.shape))))
|
||||
* mse_loss_weights
|
||||
).mean()
|
||||
|
||||
# Gather the losses across all processes for logging (if we use distributed training).
|
||||
avg_loss = accelerator.gather(
|
||||
loss.repeat(cfg.data.train_bs)).mean()
|
||||
train_loss += avg_loss.item() / cfg.solver.gradient_accumulation_steps
|
||||
|
||||
# Backpropagate
|
||||
accelerator.backward(loss)
|
||||
if accelerator.sync_gradients:
|
||||
accelerator.clip_grad_norm_(
|
||||
trainable_params,
|
||||
cfg.solver.max_grad_norm,
|
||||
)
|
||||
optimizer.step()
|
||||
lr_scheduler.step()
|
||||
optimizer.zero_grad()
|
||||
|
||||
if accelerator.sync_gradients:
|
||||
reference_control_reader.clear()
|
||||
reference_control_writer.clear()
|
||||
progress_bar.update(1)
|
||||
global_step += 1
|
||||
accelerator.log({"train_loss": train_loss}, step=global_step)
|
||||
train_loss = 0.0
|
||||
|
||||
if global_step % cfg.val.validation_steps == 0:
|
||||
if accelerator.is_main_process:
|
||||
generator = torch.Generator(device=accelerator.device)
|
||||
generator.manual_seed(cfg.seed)
|
||||
|
||||
log_validation(
|
||||
accelerator=accelerator,
|
||||
vae=vae,
|
||||
net=net,
|
||||
scheduler=val_noise_scheduler,
|
||||
width=cfg.data.train_width,
|
||||
height=cfg.data.train_height,
|
||||
clip_length=cfg.data.n_sample_frames,
|
||||
cfg=cfg,
|
||||
save_dir=validation_dir,
|
||||
global_step=global_step,
|
||||
times=cfg.single_inference_times if cfg.single_inference_times is not None else None,
|
||||
val_dataset=val_dataset,
|
||||
weight_dtype=weight_dtype
|
||||
)
|
||||
|
||||
logs = {
|
||||
"step_loss": loss.detach().item(),
|
||||
"lr": lr_scheduler.get_last_lr()[0],
|
||||
}
|
||||
|
||||
progress_bar.set_postfix(**logs)
|
||||
|
||||
writer.add_scalar("loss", loss.detach().item(), global_step)
|
||||
|
||||
if (
|
||||
global_step % cfg.checkpointing_steps == 0
|
||||
or global_step == cfg.solver.max_train_steps
|
||||
):
|
||||
# save model
|
||||
save_path = os.path.join(
|
||||
checkpoint_dir, f"checkpoint-{global_step}")
|
||||
if accelerator.is_main_process:
|
||||
delete_additional_ckpt(checkpoint_dir, 0)
|
||||
accelerator.wait_for_everyone()
|
||||
accelerator.save_state(save_path)
|
||||
|
||||
# save model weight
|
||||
unwrap_net = accelerator.unwrap_model(net)
|
||||
if accelerator.is_main_process:
|
||||
save_checkpoint(
|
||||
unwrap_net.reference_unet,
|
||||
module_dir,
|
||||
"reference_unet",
|
||||
global_step,
|
||||
total_limit=1,
|
||||
)
|
||||
save_checkpoint(
|
||||
unwrap_net.denoising_unet,
|
||||
module_dir,
|
||||
"denoising_unet",
|
||||
global_step,
|
||||
total_limit=1,
|
||||
)
|
||||
|
||||
if global_step >= cfg.solver.max_train_steps:
|
||||
break
|
||||
|
||||
# Create the pipeline using the trained modules and save it.
|
||||
accelerator.wait_for_everyone()
|
||||
accelerator.end_training()
|
||||
|
||||
|
||||
def load_config(config_path: str) -> dict:
|
||||
"""
|
||||
Loads the configuration file.
|
||||
|
||||
Args:
|
||||
config_path (str): Path to the configuration file.
|
||||
|
||||
Returns:
|
||||
dict: The configuration dictionary.
|
||||
"""
|
||||
|
||||
if config_path.endswith(".yaml"):
|
||||
return OmegaConf.load(config_path)
|
||||
if config_path.endswith(".py"):
|
||||
return import_filename(config_path).cfg
|
||||
raise ValueError("Unsupported format for config file")
|
||||
|
||||
|
||||
if __name__ == "__main__":
|
||||
parser = argparse.ArgumentParser()
|
||||
parser.add_argument(
|
||||
"--config", type=str, default="./configs/train/stage2.yaml"
|
||||
)
|
||||
args = parser.parse_args()
|
||||
|
||||
try:
|
||||
config = load_config(args.config)
|
||||
train_stage2_process(config)
|
||||
except Exception as e:
|
||||
logging.error("Failed to execute the training process: %s", e)
|
||||
@@ -0,0 +1 @@
|
||||
"""ComfyUI adapter utilities for HighSync."""
|
||||
@@ -0,0 +1,79 @@
|
||||
from pathlib import Path
|
||||
|
||||
from .file_utils import normalize_path, resolve_path
|
||||
|
||||
|
||||
CORE_FILES = (
|
||||
"denoising_unet-500.pth",
|
||||
"reference_unet-500.pth",
|
||||
"audio_processor/whisper_tiny.pt",
|
||||
)
|
||||
|
||||
ALLOW_PATTERNS = (
|
||||
"denoising_unet-500.pth",
|
||||
"reference_unet-500.pth",
|
||||
"audio_processor/whisper_tiny.pt",
|
||||
"sd-vae-ft-mse/**",
|
||||
"sd-image-variations-diffusers/**",
|
||||
)
|
||||
|
||||
|
||||
def _directory_has_model_files(path):
|
||||
if not path.is_dir():
|
||||
return False
|
||||
return any(candidate.is_file() for candidate in path.rglob("*"))
|
||||
|
||||
|
||||
def validate_required_models(models_dir):
|
||||
root = resolve_path(models_dir)
|
||||
missing = []
|
||||
|
||||
for rel_path in CORE_FILES:
|
||||
if not (root / rel_path).is_file():
|
||||
missing.append(rel_path)
|
||||
|
||||
vae_dir = root / "sd-vae-ft-mse"
|
||||
if not _directory_has_model_files(vae_dir) or not (vae_dir / "config.json").is_file():
|
||||
missing.append("sd-vae-ft-mse/config.json")
|
||||
|
||||
base_dir = root / "sd-image-variations-diffusers"
|
||||
if not _directory_has_model_files(base_dir) or not (base_dir / "unet" / "config.json").is_file():
|
||||
missing.append("sd-image-variations-diffusers/unet/config.json")
|
||||
|
||||
if missing:
|
||||
missing_text = ", ".join(missing)
|
||||
raise FileNotFoundError(
|
||||
f"HighSync model files are missing in {root}: {missing_text}. "
|
||||
"Run the HighSync Download Models node with download=True, or download "
|
||||
"https://huggingface.co/saeed-5959/high_sync into this folder."
|
||||
)
|
||||
|
||||
return root
|
||||
|
||||
|
||||
def download_or_check_models(models_dir, repo_id="saeed-5959/high_sync", download=False):
|
||||
root = Path(normalize_path(models_dir))
|
||||
|
||||
if download:
|
||||
try:
|
||||
from huggingface_hub import snapshot_download
|
||||
except ImportError as exc:
|
||||
raise RuntimeError(
|
||||
"huggingface_hub is required for HighSync Download Models. "
|
||||
"Install this wrapper's requirements.txt first."
|
||||
) from exc
|
||||
|
||||
root.mkdir(parents=True, exist_ok=True)
|
||||
snapshot_download(
|
||||
repo_id=repo_id,
|
||||
local_dir=str(root),
|
||||
allow_patterns=list(ALLOW_PATTERNS),
|
||||
repo_type="model",
|
||||
)
|
||||
elif not root.exists():
|
||||
raise FileNotFoundError(
|
||||
f"models_dir does not exist: {root}. "
|
||||
"Run HighSync Download Models with download=True or download the models manually."
|
||||
)
|
||||
|
||||
return str(validate_required_models(root))
|
||||
@@ -0,0 +1,45 @@
|
||||
import shutil
|
||||
import subprocess
|
||||
|
||||
from .file_utils import resolve_path
|
||||
|
||||
|
||||
def check_ffmpeg():
|
||||
ffmpeg = shutil.which("ffmpeg") or shutil.which("ffmpeg.exe")
|
||||
if ffmpeg is None:
|
||||
raise RuntimeError("ffmpeg not found. Please install ffmpeg and add it to PATH.")
|
||||
return ffmpeg
|
||||
|
||||
|
||||
def run_ffmpeg(args, error_label):
|
||||
command = [check_ffmpeg()] + [str(arg) for arg in args]
|
||||
result = subprocess.run(command, capture_output=True, text=True)
|
||||
if result.returncode != 0:
|
||||
stderr = result.stderr.strip()
|
||||
raise RuntimeError(f"{error_label} failed with exit code {result.returncode}: {stderr}")
|
||||
|
||||
|
||||
def mux_video_audio(video_path, audio_path, output_path):
|
||||
output_file = resolve_path(output_path)
|
||||
output_file.parent.mkdir(parents=True, exist_ok=True)
|
||||
run_ffmpeg(
|
||||
[
|
||||
"-y",
|
||||
"-i",
|
||||
resolve_path(video_path),
|
||||
"-i",
|
||||
resolve_path(audio_path),
|
||||
"-map",
|
||||
"0:v:0",
|
||||
"-map",
|
||||
"1:a:0",
|
||||
"-c:v",
|
||||
"copy",
|
||||
"-c:a",
|
||||
"aac",
|
||||
"-shortest",
|
||||
output_file,
|
||||
],
|
||||
"audio/video mux",
|
||||
)
|
||||
return output_file
|
||||
@@ -0,0 +1,92 @@
|
||||
import os
|
||||
import re
|
||||
import shutil
|
||||
import uuid
|
||||
from datetime import datetime
|
||||
from pathlib import Path
|
||||
|
||||
|
||||
WINDOWS_DRIVE_RE = re.compile(r"^([A-Za-z]):[\\/](.*)$")
|
||||
|
||||
|
||||
def is_windows_drive_path(path):
|
||||
return isinstance(path, str) and WINDOWS_DRIVE_RE.match(path.strip()) is not None
|
||||
|
||||
|
||||
def normalize_path(path):
|
||||
if path is None:
|
||||
return path
|
||||
|
||||
raw_path = str(path).strip().strip('"')
|
||||
if raw_path == "":
|
||||
return raw_path
|
||||
|
||||
match = WINDOWS_DRIVE_RE.match(raw_path)
|
||||
if match and os.name != "nt":
|
||||
drive = match.group(1).lower()
|
||||
rest = match.group(2).replace("\\", "/")
|
||||
return f"/mnt/{drive}/{rest}"
|
||||
|
||||
return os.path.abspath(os.path.expanduser(raw_path))
|
||||
|
||||
|
||||
def resolve_path(path):
|
||||
return Path(normalize_path(path)).expanduser().resolve()
|
||||
|
||||
|
||||
def ensure_exists(path, label="path", is_dir=None):
|
||||
if path is None or str(path).strip() == "":
|
||||
raise FileNotFoundError(f"{label} is empty.")
|
||||
|
||||
resolved = resolve_path(path)
|
||||
if is_dir is True and not resolved.is_dir():
|
||||
raise FileNotFoundError(f"{label} does not exist: {resolved}")
|
||||
if is_dir is False and not resolved.is_file():
|
||||
raise FileNotFoundError(f"{label} does not exist: {resolved}")
|
||||
if is_dir is None and not resolved.exists():
|
||||
raise FileNotFoundError(f"{label} does not exist: {resolved}")
|
||||
return resolved
|
||||
|
||||
|
||||
def safe_copy(src, dst):
|
||||
src_path = ensure_exists(src, "source file", is_dir=False)
|
||||
dst_path = resolve_path(dst)
|
||||
dst_path.parent.mkdir(parents=True, exist_ok=True)
|
||||
if src_path.resolve() != dst_path.resolve():
|
||||
shutil.copy2(src_path, dst_path)
|
||||
return dst_path
|
||||
|
||||
|
||||
def get_comfy_output_dir():
|
||||
try:
|
||||
import folder_paths
|
||||
|
||||
return Path(folder_paths.get_output_directory()).resolve()
|
||||
except Exception:
|
||||
current = Path(__file__).resolve()
|
||||
for parent in current.parents:
|
||||
if (parent / "custom_nodes").exists() and (parent / "output").exists():
|
||||
return (parent / "output").resolve()
|
||||
return Path("/mnt/d/ComfyUI/output").resolve()
|
||||
|
||||
|
||||
def create_task_dir(prefix="highsync"):
|
||||
output_root = get_comfy_output_dir() / "highsync"
|
||||
timestamp = datetime.now().strftime("%Y%m%d_%H%M%S")
|
||||
suffix = uuid.uuid4().hex[:8]
|
||||
task_id = f"{prefix}_{timestamp}_{suffix}"
|
||||
task_dir = output_root / task_id
|
||||
|
||||
paths = {
|
||||
"task_id": task_id,
|
||||
"task_dir": task_dir,
|
||||
"media_dir": task_dir / "media",
|
||||
"logs_dir": task_dir / "logs",
|
||||
"frames_dir": task_dir / "frames",
|
||||
}
|
||||
|
||||
for key, value in paths.items():
|
||||
if key.endswith("_dir"):
|
||||
value.mkdir(parents=True, exist_ok=True)
|
||||
|
||||
return paths
|
||||
@@ -0,0 +1,224 @@
|
||||
import subprocess
|
||||
from pathlib import Path
|
||||
|
||||
import numpy as np
|
||||
import soundfile as sf
|
||||
import torch
|
||||
|
||||
from .ffmpeg_utils import check_ffmpeg
|
||||
from .file_utils import resolve_path
|
||||
|
||||
|
||||
TARGET_FPS = 25
|
||||
TARGET_AUDIO_SR = 16000
|
||||
|
||||
|
||||
def _as_image_tensor(images):
|
||||
if isinstance(images, list):
|
||||
images = torch.stack(images)
|
||||
if not isinstance(images, torch.Tensor):
|
||||
images = torch.as_tensor(images)
|
||||
if images.dim() != 4:
|
||||
raise ValueError(f"Expected IMAGE tensor with shape [frames, height, width, channels], got {tuple(images.shape)}")
|
||||
if images.shape[-1] < 3:
|
||||
raise ValueError(f"Expected IMAGE tensor with at least 3 channels, got {tuple(images.shape)}")
|
||||
return images.detach().cpu()
|
||||
|
||||
|
||||
def images_to_uint8_rgb(images):
|
||||
image_tensor = _as_image_tensor(images)[..., :3]
|
||||
if image_tensor.dtype != torch.uint8:
|
||||
image_tensor = image_tensor.float()
|
||||
if image_tensor.numel() > 0 and float(image_tensor.max()) <= 1.0:
|
||||
image_tensor = image_tensor * 255.0
|
||||
image_tensor = image_tensor.clamp(0, 255).byte()
|
||||
return np.ascontiguousarray(image_tensor.numpy())
|
||||
|
||||
|
||||
def frames_uint8_rgb_to_image_tensor(frames):
|
||||
if len(frames) == 0:
|
||||
raise ValueError("No output frames were generated.")
|
||||
array = np.ascontiguousarray(np.stack(frames).astype(np.float32) / 255.0)
|
||||
return torch.from_numpy(array)
|
||||
|
||||
|
||||
def resample_frames_to_fps(frames_rgb, input_fps, output_fps=TARGET_FPS, max_frames=0):
|
||||
input_fps = int(input_fps)
|
||||
output_fps = int(output_fps)
|
||||
if input_fps <= 0 or output_fps <= 0:
|
||||
raise ValueError("input_fps and output_fps must be positive.")
|
||||
if len(frames_rgb) == 0:
|
||||
raise ValueError("No IMAGE frames were provided.")
|
||||
|
||||
if input_fps == output_fps:
|
||||
sampled = frames_rgb.copy()
|
||||
else:
|
||||
duration = len(frames_rgb) / float(input_fps)
|
||||
target_count = max(1, int(round(duration * output_fps)))
|
||||
indices = np.floor(np.arange(target_count) * input_fps / output_fps).astype(np.int64)
|
||||
indices = np.clip(indices, 0, len(frames_rgb) - 1)
|
||||
sampled = frames_rgb[indices]
|
||||
|
||||
max_frames = int(max_frames or 0)
|
||||
if max_frames > 0:
|
||||
sampled = sampled[:max_frames]
|
||||
|
||||
return np.ascontiguousarray(sampled)
|
||||
|
||||
|
||||
def _audio_waveform_to_tensor(audio):
|
||||
if not isinstance(audio, dict) or "waveform" not in audio or "sample_rate" not in audio:
|
||||
raise ValueError("AUDIO input must be a dict with 'waveform' and 'sample_rate'.")
|
||||
|
||||
waveform = audio["waveform"]
|
||||
if not isinstance(waveform, torch.Tensor):
|
||||
waveform = torch.as_tensor(waveform)
|
||||
waveform = waveform.detach().cpu().float()
|
||||
|
||||
if waveform.dim() == 3:
|
||||
waveform = waveform[0]
|
||||
elif waveform.dim() == 1:
|
||||
waveform = waveform.unsqueeze(0)
|
||||
elif waveform.dim() != 2:
|
||||
raise ValueError(
|
||||
"Expected AUDIO waveform with shape [batch, channels, samples], "
|
||||
f"[channels, samples], or [samples], got {tuple(waveform.shape)}"
|
||||
)
|
||||
|
||||
return waveform.contiguous(), int(audio["sample_rate"])
|
||||
|
||||
|
||||
def trim_or_pad_audio(audio, duration_sec=None, pad_to_duration=False):
|
||||
waveform, sample_rate = _audio_waveform_to_tensor(audio)
|
||||
|
||||
if duration_sec is not None and duration_sec > 0:
|
||||
target_samples = max(1, int(round(float(duration_sec) * sample_rate)))
|
||||
if waveform.shape[-1] > target_samples:
|
||||
waveform = waveform[:, :target_samples]
|
||||
elif pad_to_duration and waveform.shape[-1] < target_samples:
|
||||
pad = torch.zeros((waveform.shape[0], target_samples - waveform.shape[-1]), dtype=waveform.dtype)
|
||||
waveform = torch.cat([waveform, pad], dim=-1)
|
||||
|
||||
return {"waveform": waveform.unsqueeze(0).contiguous(), "sample_rate": sample_rate}
|
||||
|
||||
|
||||
def _resample_waveform(waveform, sample_rate, target_sample_rate):
|
||||
if int(sample_rate) == int(target_sample_rate):
|
||||
return waveform
|
||||
|
||||
try:
|
||||
import torchaudio
|
||||
|
||||
resampler = torchaudio.transforms.Resample(
|
||||
orig_freq=int(sample_rate),
|
||||
new_freq=int(target_sample_rate),
|
||||
)
|
||||
return resampler(waveform)
|
||||
except Exception as exc:
|
||||
raise RuntimeError(
|
||||
"Failed to resample audio with torchaudio. Install torchaudio or provide 16kHz audio."
|
||||
) from exc
|
||||
|
||||
|
||||
def save_audio_to_wav(audio, output_path, duration_sec=None, pad_to_duration=False, mono=False, sample_rate=None):
|
||||
adjusted = trim_or_pad_audio(audio, duration_sec=duration_sec, pad_to_duration=pad_to_duration)
|
||||
waveform = adjusted["waveform"].squeeze(0).clamp(-1.0, 1.0)
|
||||
source_sample_rate = int(adjusted["sample_rate"])
|
||||
|
||||
if mono and waveform.shape[0] > 1:
|
||||
waveform = waveform.mean(dim=0, keepdim=True)
|
||||
|
||||
target_sample_rate = int(sample_rate or source_sample_rate)
|
||||
waveform = _resample_waveform(waveform, source_sample_rate, target_sample_rate)
|
||||
|
||||
output_file = resolve_path(output_path)
|
||||
output_file.parent.mkdir(parents=True, exist_ok=True)
|
||||
audio_np = waveform.transpose(0, 1).numpy()
|
||||
sf.write(str(output_file), audio_np, target_sample_rate, subtype="PCM_16")
|
||||
return output_file
|
||||
|
||||
|
||||
def encode_rgb_frames_to_mp4(frames_rgb, output_path, fps=TARGET_FPS):
|
||||
output_file = resolve_path(output_path)
|
||||
output_file.parent.mkdir(parents=True, exist_ok=True)
|
||||
|
||||
frames = np.ascontiguousarray(frames_rgb)
|
||||
if frames.ndim != 4 or frames.shape[-1] != 3:
|
||||
raise ValueError(f"Expected RGB frames with shape [frames, height, width, 3], got {frames.shape}")
|
||||
|
||||
frame_count, height, width, _ = frames.shape
|
||||
if frame_count <= 0:
|
||||
raise ValueError("No frames were provided to encode.")
|
||||
|
||||
command = [
|
||||
check_ffmpeg(),
|
||||
"-y",
|
||||
"-f",
|
||||
"rawvideo",
|
||||
"-vcodec",
|
||||
"rawvideo",
|
||||
"-s",
|
||||
f"{width}x{height}",
|
||||
"-pix_fmt",
|
||||
"rgb24",
|
||||
"-r",
|
||||
str(int(fps)),
|
||||
"-i",
|
||||
"-",
|
||||
"-an",
|
||||
"-c:v",
|
||||
"libx264",
|
||||
"-pix_fmt",
|
||||
"yuv420p",
|
||||
"-crf",
|
||||
"18",
|
||||
str(output_file),
|
||||
]
|
||||
|
||||
process = subprocess.Popen(command, stdin=subprocess.PIPE, stdout=subprocess.PIPE, stderr=subprocess.PIPE)
|
||||
try:
|
||||
for frame in frames:
|
||||
process.stdin.write(np.ascontiguousarray(frame).tobytes())
|
||||
_, stderr = process.communicate()
|
||||
except BrokenPipeError:
|
||||
_, stderr = process.communicate()
|
||||
finally:
|
||||
if process.stdin and not process.stdin.closed:
|
||||
process.stdin.close()
|
||||
|
||||
if process.returncode != 0:
|
||||
message = stderr.decode("utf-8", errors="replace").strip()
|
||||
raise RuntimeError(f"Failed to encode HighSync output video: {message}")
|
||||
|
||||
return output_file
|
||||
|
||||
|
||||
def denoise_audio_file(input_audio_path, output_audio_path, device, cache):
|
||||
try:
|
||||
from denoiser import pretrained
|
||||
from denoiser.dsp import convert_audio
|
||||
except ImportError as exc:
|
||||
raise RuntimeError(
|
||||
"denoiser and torchaudio are required when denoise_audio=True. "
|
||||
"Install this wrapper's requirements.txt or disable denoise_audio."
|
||||
) from exc
|
||||
|
||||
model = cache.get("denoiser_model")
|
||||
model_device = torch.device(device)
|
||||
if model is None:
|
||||
model = pretrained.dns64().to(model_device)
|
||||
model.eval()
|
||||
cache["denoiser_model"] = model
|
||||
|
||||
audio_np, sample_rate = sf.read(str(Path(input_audio_path)), dtype="float32", always_2d=True)
|
||||
wav = torch.from_numpy(audio_np.T.copy())
|
||||
wav = convert_audio(wav.to(model_device), sample_rate, model.sample_rate, model.chin)
|
||||
|
||||
with torch.no_grad():
|
||||
denoised = model(wav[None])[0]
|
||||
|
||||
output_file = resolve_path(output_audio_path)
|
||||
output_file.parent.mkdir(parents=True, exist_ok=True)
|
||||
audio_np = denoised.detach().cpu().clamp(-1.0, 1.0).numpy()[0]
|
||||
sf.write(str(output_file), audio_np, int(model.sample_rate), subtype="PCM_16")
|
||||
return output_file
|
||||
@@ -0,0 +1,148 @@
|
||||
from pathlib import Path
|
||||
|
||||
from .download_utils import validate_required_models
|
||||
|
||||
|
||||
_MODEL_CACHE = {}
|
||||
|
||||
|
||||
def _get_repo_root():
|
||||
return Path(__file__).resolve().parents[1]
|
||||
|
||||
|
||||
def _torch_load(path, torch):
|
||||
try:
|
||||
return torch.load(path, map_location="cpu", weights_only=True)
|
||||
except TypeError:
|
||||
return torch.load(path, map_location="cpu")
|
||||
|
||||
|
||||
def _get_torch_device(torch):
|
||||
try:
|
||||
import comfy.model_management as model_management
|
||||
|
||||
return model_management.get_torch_device()
|
||||
except Exception:
|
||||
return torch.device("cuda" if torch.cuda.is_available() else "cpu")
|
||||
|
||||
|
||||
def _weight_dtype(torch, weight_dtype):
|
||||
dtype_text = str(weight_dtype or "fp16").lower()
|
||||
if dtype_text == "fp16":
|
||||
return torch.float16
|
||||
if dtype_text == "bf16":
|
||||
return torch.bfloat16
|
||||
if dtype_text == "fp32":
|
||||
return torch.float32
|
||||
raise ValueError(f"Unsupported HighSync weight_dtype: {weight_dtype}")
|
||||
|
||||
|
||||
def _load_config(models_root, weight_dtype):
|
||||
try:
|
||||
from omegaconf import OmegaConf
|
||||
except ImportError as exc:
|
||||
raise RuntimeError("omegaconf is required to load HighSync config. Install requirements.txt first.") from exc
|
||||
|
||||
config_path = _get_repo_root() / "high_sync" / "configs" / "inference" / "default.yaml"
|
||||
config = OmegaConf.load(str(config_path))
|
||||
|
||||
config.weight_dtype = str(weight_dtype)
|
||||
config.denoising_unet_path = str(models_root / "denoising_unet-500.pth")
|
||||
config.reference_unet_path = str(models_root / "reference_unet-500.pth")
|
||||
config.base_model_path = str(models_root / "sd-image-variations-diffusers")
|
||||
config.audio_model_path = str(models_root / "audio_processor" / "whisper_tiny.pt")
|
||||
config.vae.model_path = str(models_root / "sd-vae-ft-mse")
|
||||
config.mask_path = str(_get_repo_root() / "high_sync" / "masks" / "mask_64.png")
|
||||
config.mask_512_path = str(_get_repo_root() / "high_sync" / "masks" / "mask_512.png")
|
||||
|
||||
return config
|
||||
|
||||
|
||||
def _build_scheduler(config):
|
||||
try:
|
||||
from diffusers import DDIMScheduler
|
||||
from omegaconf import OmegaConf
|
||||
except ImportError as exc:
|
||||
raise RuntimeError("diffusers and omegaconf are required for HighSync. Install requirements.txt first.") from exc
|
||||
|
||||
sched_kwargs = OmegaConf.to_container(config.noise_scheduler_kwargs, resolve=True)
|
||||
if bool(config.enable_zero_snr):
|
||||
sched_kwargs.update(
|
||||
rescale_betas_zero_snr=True,
|
||||
timestep_spacing="trailing",
|
||||
prediction_type="v_prediction",
|
||||
)
|
||||
return DDIMScheduler(**sched_kwargs)
|
||||
|
||||
|
||||
def load_highsync_model(models_dir, weight_dtype="fp16"):
|
||||
try:
|
||||
import torch
|
||||
from diffusers import AutoencoderKL
|
||||
except ImportError as exc:
|
||||
raise RuntimeError(
|
||||
"torch and diffusers are required to load HighSync. Install this wrapper's requirements.txt first."
|
||||
) from exc
|
||||
|
||||
models_root = validate_required_models(models_dir)
|
||||
device = _get_torch_device(torch)
|
||||
dtype = _weight_dtype(torch, weight_dtype)
|
||||
cache_key = (str(models_root), str(weight_dtype).lower(), str(device))
|
||||
|
||||
if cache_key in _MODEL_CACHE:
|
||||
return _MODEL_CACHE[cache_key]
|
||||
|
||||
config = _load_config(models_root, weight_dtype)
|
||||
|
||||
from omegaconf import OmegaConf
|
||||
|
||||
from high_sync.image_processor import ImageProcessor
|
||||
from high_sync.src.models.unet_2d_condition import UNet2DConditionModel
|
||||
from high_sync.src.models.unet_3d_echo import EchoUNet3DConditionModel
|
||||
from high_sync.src.models.whisper.audio2feature import load_audio_model
|
||||
from high_sync.src.pipelines.face_animate import FaceAnimatePipeline
|
||||
|
||||
scheduler = _build_scheduler(config)
|
||||
|
||||
vae = AutoencoderKL.from_pretrained(str(config.vae.model_path))
|
||||
reference_unet = UNet2DConditionModel.from_pretrained(str(config.base_model_path), subfolder="unet")
|
||||
denoising_unet = EchoUNet3DConditionModel.from_pretrained_2d(
|
||||
str(config.base_model_path),
|
||||
"",
|
||||
subfolder="unet",
|
||||
unet_additional_kwargs=OmegaConf.to_container(config.unet_additional_kwargs, resolve=True),
|
||||
)
|
||||
|
||||
reference_unet.load_state_dict(_torch_load(str(config.reference_unet_path), torch))
|
||||
denoising_unet.load_state_dict(_torch_load(str(config.denoising_unet_path), torch))
|
||||
|
||||
vae.requires_grad_(False)
|
||||
reference_unet.requires_grad_(False)
|
||||
denoising_unet.requires_grad_(False)
|
||||
|
||||
vae.eval()
|
||||
reference_unet.eval()
|
||||
denoising_unet.eval()
|
||||
|
||||
pipeline = FaceAnimatePipeline(
|
||||
vae=vae,
|
||||
reference_unet=reference_unet,
|
||||
denoising_unet=denoising_unet,
|
||||
scheduler=scheduler,
|
||||
)
|
||||
pipeline.to(device=device, dtype=dtype)
|
||||
|
||||
audio_processor = load_audio_model(model_path=str(config.audio_model_path), device=device)
|
||||
|
||||
model = {
|
||||
"models_dir": str(models_root),
|
||||
"device": device,
|
||||
"weight_dtype": dtype,
|
||||
"config": config,
|
||||
"pipeline": pipeline,
|
||||
"audio_processor": audio_processor,
|
||||
"image_processor_cls": ImageProcessor,
|
||||
"denoiser_model": None,
|
||||
}
|
||||
_MODEL_CACHE[cache_key] = model
|
||||
return model
|
||||
@@ -0,0 +1,323 @@
|
||||
import shutil
|
||||
|
||||
import cv2
|
||||
import numpy as np
|
||||
import torch
|
||||
|
||||
from .ffmpeg_utils import mux_video_audio
|
||||
from .file_utils import create_task_dir
|
||||
from .media_utils import (
|
||||
TARGET_AUDIO_SR,
|
||||
TARGET_FPS,
|
||||
denoise_audio_file,
|
||||
encode_rgb_frames_to_mp4,
|
||||
frames_uint8_rgb_to_image_tensor,
|
||||
images_to_uint8_rgb,
|
||||
resample_frames_to_fps,
|
||||
save_audio_to_wav,
|
||||
trim_or_pad_audio,
|
||||
)
|
||||
|
||||
|
||||
def _maybe_throw_if_interrupted():
|
||||
try:
|
||||
import comfy.model_management as model_management
|
||||
except Exception:
|
||||
return
|
||||
model_management.throw_exception_if_processing_interrupted()
|
||||
|
||||
|
||||
def _empty_cache(device):
|
||||
if isinstance(device, torch.device) and device.type == "cuda" and torch.cuda.is_available():
|
||||
torch.cuda.empty_cache()
|
||||
try:
|
||||
import comfy.model_management as model_management
|
||||
|
||||
model_management.soft_empty_cache()
|
||||
except Exception:
|
||||
return
|
||||
|
||||
|
||||
def create_mask(mask_path):
|
||||
raw_mask = cv2.imread(str(mask_path))
|
||||
if raw_mask is None:
|
||||
raise FileNotFoundError(f"HighSync mask could not be read: {mask_path}")
|
||||
mask = 1 - (raw_mask / 255.0)
|
||||
|
||||
h, w, c = mask.shape
|
||||
new_mask = np.zeros((h, w, c + 1), dtype=np.float32)
|
||||
new_mask[:, :, :c] = mask
|
||||
new_mask[:, :, c] = mask[:, :, c - 1]
|
||||
return torch.FloatTensor(new_mask).permute(2, 0, 1).unsqueeze(1)
|
||||
|
||||
|
||||
def create_mask_alpha(width, height, coordinate):
|
||||
left, top, right, bottom = coordinate
|
||||
mask = np.ones((height, width, 3), dtype=np.float32)
|
||||
mask[top:bottom, left:right] = 0
|
||||
|
||||
pad = min(int((right - left) / 15), int((bottom - top) / 15))
|
||||
if pad <= 0:
|
||||
return mask
|
||||
|
||||
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 _prepare_audio_files(task, audio, duration_sec, trim_audio_to_video, denoise_audio, device, model_cache):
|
||||
audio_for_video = trim_or_pad_audio(
|
||||
audio,
|
||||
duration_sec=duration_sec if trim_audio_to_video else None,
|
||||
pad_to_duration=bool(trim_audio_to_video),
|
||||
)
|
||||
|
||||
mux_audio_path = save_audio_to_wav(
|
||||
audio_for_video,
|
||||
task["media_dir"] / "mux_audio.wav",
|
||||
)
|
||||
feature_audio_path = save_audio_to_wav(
|
||||
audio_for_video,
|
||||
task["media_dir"] / "feature_audio_16k.wav",
|
||||
mono=True,
|
||||
sample_rate=TARGET_AUDIO_SR,
|
||||
)
|
||||
|
||||
if denoise_audio:
|
||||
feature_audio_path = denoise_audio_file(
|
||||
feature_audio_path,
|
||||
task["media_dir"] / "feature_audio_16k_denoised.wav",
|
||||
device=device,
|
||||
cache=model_cache,
|
||||
)
|
||||
|
||||
return audio_for_video, mux_audio_path, feature_audio_path
|
||||
|
||||
|
||||
def _ensure_highsync_model(highsync_model):
|
||||
required = ("config", "pipeline", "audio_processor", "image_processor_cls", "device", "weight_dtype")
|
||||
if not isinstance(highsync_model, dict) or any(key not in highsync_model for key in required):
|
||||
raise ValueError("highsync_model must come from the HighSync Model Loader node.")
|
||||
|
||||
|
||||
def _append_original_frame(output_frames_rgb, frame_bgr):
|
||||
output_frames_rgb.append(cv2.cvtColor(frame_bgr, cv2.COLOR_BGR2RGB))
|
||||
|
||||
|
||||
def _run_generation_chunk(
|
||||
pipeline,
|
||||
image_processor,
|
||||
config,
|
||||
audio_fea_final,
|
||||
frames_video,
|
||||
output_frames_rgb,
|
||||
count_img,
|
||||
generator,
|
||||
device,
|
||||
weight_dtype,
|
||||
inference_steps,
|
||||
cfg_scale,
|
||||
):
|
||||
source_image_pixels_all, bbox_all, frame_list_all = image_processor.preprocess_frames_jump(frames_video)
|
||||
|
||||
img_size = (int(config.data.source_image.width), int(config.data.source_image.height))
|
||||
sample_frames = int(config.data.n_sample_frames)
|
||||
mask_path = str(config.mask_path)
|
||||
|
||||
for source_image_pixels_list, bbox_list, frame_list in zip(source_image_pixels_all, bbox_all, frame_list_all):
|
||||
_maybe_throw_if_interrupted()
|
||||
|
||||
if not frame_list:
|
||||
continue
|
||||
|
||||
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 = sample_frames - 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).unsqueeze(0).to(device=device, dtype=weight_dtype)
|
||||
_empty_cache(device)
|
||||
|
||||
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=int(inference_steps),
|
||||
guidance_scale=float(cfg_scale),
|
||||
generator=generator,
|
||||
)
|
||||
|
||||
_empty_cache(device)
|
||||
|
||||
images_out = pipeline_output.videos.squeeze(0).permute(1, 2, 3, 0).cpu().numpy()[:main_length]
|
||||
images_out = np.clip(images_out * 255, 0, 255).astype(np.uint8)
|
||||
|
||||
for image_out, bbox, frame in zip(images_out, bbox_list, frame_list):
|
||||
frame_h, frame_w, _ = frame.shape
|
||||
if bbox == (0, 0, 0, 0):
|
||||
_append_original_frame(output_frames_rgb, frame)
|
||||
else:
|
||||
image_out_bgr = 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_bgr,
|
||||
(bbox[2] - bbox[0], bbox[3] - bbox[1]),
|
||||
)
|
||||
mask = create_mask_alpha(frame_w, frame_h, bbox)
|
||||
blended = (frame * mask + frame_in_out * (1 - mask)).astype(np.uint8)
|
||||
_append_original_frame(output_frames_rgb, blended)
|
||||
|
||||
count_img += 1
|
||||
else:
|
||||
for frame in frame_list:
|
||||
_append_original_frame(output_frames_rgb, frame)
|
||||
count_img += 1
|
||||
|
||||
return count_img
|
||||
|
||||
|
||||
def run_highsync(
|
||||
highsync_model,
|
||||
images,
|
||||
audio,
|
||||
input_fps=25,
|
||||
max_frames=0,
|
||||
inference_steps=None,
|
||||
cfg_scale=None,
|
||||
seed=42,
|
||||
preprocess_chunk_frames=1500,
|
||||
denoise_audio=True,
|
||||
trim_audio_to_video=True,
|
||||
keep_intermediate_outputs=True,
|
||||
):
|
||||
_ensure_highsync_model(highsync_model)
|
||||
|
||||
config = highsync_model["config"]
|
||||
pipeline = highsync_model["pipeline"]
|
||||
audio_processor = highsync_model["audio_processor"]
|
||||
image_processor = highsync_model["image_processor_cls"]((int(config.data.source_image.width), int(config.data.source_image.height)))
|
||||
device = torch.device(highsync_model["device"])
|
||||
weight_dtype = highsync_model["weight_dtype"]
|
||||
|
||||
inference_steps = int(inference_steps if inference_steps is not None else config.inference_steps)
|
||||
cfg_scale = float(cfg_scale if cfg_scale is not None else config.cfg_scale)
|
||||
preprocess_chunk_frames = max(12, int(preprocess_chunk_frames))
|
||||
|
||||
task = create_task_dir()
|
||||
|
||||
frames_rgb = images_to_uint8_rgb(images)
|
||||
frames_rgb = resample_frames_to_fps(frames_rgb, input_fps=input_fps, output_fps=TARGET_FPS, max_frames=max_frames)
|
||||
frames_bgr = [cv2.cvtColor(frame, cv2.COLOR_RGB2BGR) for frame in frames_rgb]
|
||||
initial_duration_sec = len(frames_bgr) / float(TARGET_FPS)
|
||||
|
||||
_, _, feature_audio_path = _prepare_audio_files(
|
||||
task=task,
|
||||
audio=audio,
|
||||
duration_sec=initial_duration_sec,
|
||||
trim_audio_to_video=trim_audio_to_video,
|
||||
denoise_audio=denoise_audio,
|
||||
device=device,
|
||||
model_cache=highsync_model,
|
||||
)
|
||||
|
||||
whisper_feature = audio_processor.audio2feat(str(feature_audio_path))
|
||||
whisper_chunks = audio_processor.feature2chunks(feature_array=whisper_feature, fps=TARGET_FPS)
|
||||
audio_frame_num = int(whisper_chunks.shape[0])
|
||||
video_length = min(len(frames_bgr), audio_frame_num)
|
||||
if video_length <= 0:
|
||||
raise RuntimeError("HighSync could not align any video frames with the provided audio.")
|
||||
|
||||
frames_bgr = frames_bgr[:video_length]
|
||||
audio_fea_final = torch.as_tensor(whisper_chunks[:video_length], dtype=pipeline.vae.dtype, device=pipeline.vae.device)
|
||||
audio_fea_final = audio_fea_final.unsqueeze(0)
|
||||
|
||||
final_duration_sec = video_length / float(TARGET_FPS)
|
||||
output_audio, mux_audio_path, _ = _prepare_audio_files(
|
||||
task=task,
|
||||
audio=audio,
|
||||
duration_sec=final_duration_sec,
|
||||
trim_audio_to_video=trim_audio_to_video,
|
||||
denoise_audio=False,
|
||||
device=device,
|
||||
model_cache=highsync_model,
|
||||
)
|
||||
|
||||
torch.manual_seed(int(seed))
|
||||
generator = torch.manual_seed(int(seed))
|
||||
|
||||
output_frames_rgb = []
|
||||
frames_video = []
|
||||
count_motion = 0
|
||||
count_img = 0
|
||||
|
||||
while True:
|
||||
_maybe_throw_if_interrupted()
|
||||
|
||||
if count_motion >= video_length:
|
||||
finish = True
|
||||
else:
|
||||
frames_video.append(frames_bgr[count_motion])
|
||||
count_motion += 1
|
||||
finish = False
|
||||
|
||||
if len(frames_video) == 0:
|
||||
break
|
||||
|
||||
if (count_motion % preprocess_chunk_frames == 0) or finish:
|
||||
print(f"[HighSync] processing frames {count_img}-{count_motion} / {video_length}")
|
||||
count_img = _run_generation_chunk(
|
||||
pipeline=pipeline,
|
||||
image_processor=image_processor,
|
||||
config=config,
|
||||
audio_fea_final=audio_fea_final,
|
||||
frames_video=frames_video,
|
||||
output_frames_rgb=output_frames_rgb,
|
||||
count_img=count_img,
|
||||
generator=generator,
|
||||
device=device,
|
||||
weight_dtype=weight_dtype,
|
||||
inference_steps=inference_steps,
|
||||
cfg_scale=cfg_scale,
|
||||
)
|
||||
frames_video = []
|
||||
|
||||
if finish:
|
||||
break
|
||||
|
||||
if len(output_frames_rgb) == 0:
|
||||
raise RuntimeError("HighSync finished without producing output frames.")
|
||||
|
||||
output_frames = np.ascontiguousarray(np.stack(output_frames_rgb))
|
||||
no_audio_path = encode_rgb_frames_to_mp4(output_frames, task["media_dir"] / "highsync_no_audio.mp4", fps=TARGET_FPS)
|
||||
final_video_path = mux_video_audio(no_audio_path, mux_audio_path, task["task_dir"] / "final.mp4")
|
||||
|
||||
output_images = frames_uint8_rgb_to_image_tensor(output_frames)
|
||||
|
||||
if not keep_intermediate_outputs:
|
||||
for path in (task["media_dir"], task["frames_dir"], task["logs_dir"]):
|
||||
shutil.rmtree(path, ignore_errors=True)
|
||||
|
||||
return output_images, output_audio, TARGET_FPS, str(final_video_path.resolve())
|
||||
@@ -0,0 +1,169 @@
|
||||
try:
|
||||
from .highsync_wrapper.download_utils import download_or_check_models
|
||||
except ImportError:
|
||||
from highsync_wrapper.download_utils import download_or_check_models
|
||||
|
||||
|
||||
DEFAULT_MODELS_DIR = "D:/ComfyUI/models/highsync"
|
||||
|
||||
|
||||
class HighSyncDownloadModels:
|
||||
@classmethod
|
||||
def INPUT_TYPES(cls):
|
||||
return {
|
||||
"required": {
|
||||
"download": ("BOOLEAN", {"default": False}),
|
||||
}
|
||||
}
|
||||
|
||||
RETURN_TYPES = ("STRING",)
|
||||
RETURN_NAMES = ("models_dir",)
|
||||
FUNCTION = "run"
|
||||
CATEGORY = "video/lipsync"
|
||||
|
||||
def run(self, download):
|
||||
models_dir = download_or_check_models(
|
||||
models_dir=DEFAULT_MODELS_DIR,
|
||||
repo_id="saeed-5959/high_sync",
|
||||
download=download,
|
||||
)
|
||||
return (models_dir,)
|
||||
|
||||
|
||||
class HighSyncModelLoader:
|
||||
@classmethod
|
||||
def INPUT_TYPES(cls):
|
||||
return {
|
||||
"required": {
|
||||
"models_dir": ("STRING", {"default": DEFAULT_MODELS_DIR, "multiline": False}),
|
||||
"weight_dtype": (["fp16", "bf16", "fp32"], {"default": "fp16"}),
|
||||
}
|
||||
}
|
||||
|
||||
RETURN_TYPES = ("HIGHSYNC_MODEL",)
|
||||
RETURN_NAMES = ("highsync_model",)
|
||||
FUNCTION = "load"
|
||||
CATEGORY = "video/lipsync"
|
||||
|
||||
def load(self, models_dir, weight_dtype):
|
||||
try:
|
||||
from .highsync_wrapper.model_utils import load_highsync_model
|
||||
except ImportError:
|
||||
from highsync_wrapper.model_utils import load_highsync_model
|
||||
|
||||
model = load_highsync_model(models_dir=models_dir, weight_dtype=weight_dtype)
|
||||
return (model,)
|
||||
|
||||
|
||||
class HighSyncLipSync:
|
||||
@classmethod
|
||||
def INPUT_TYPES(cls):
|
||||
return {
|
||||
"required": {
|
||||
"highsync_model": ("HIGHSYNC_MODEL",),
|
||||
"images": ("IMAGE",),
|
||||
"audio": ("AUDIO",),
|
||||
"input_fps": ("INT", {"default": 25, "min": 1, "max": 120, "step": 1}),
|
||||
"max_frames": ("INT", {"default": 0, "min": 0, "max": 100000, "step": 1}),
|
||||
}
|
||||
}
|
||||
|
||||
RETURN_TYPES = ("IMAGE", "AUDIO", "INT", "STRING")
|
||||
RETURN_NAMES = ("images", "audio", "frame_rate", "output_video_path")
|
||||
FUNCTION = "run"
|
||||
CATEGORY = "video/lipsync"
|
||||
|
||||
def run(self, highsync_model, images, audio, input_fps, max_frames):
|
||||
try:
|
||||
from .highsync_wrapper.runner import run_highsync
|
||||
except ImportError:
|
||||
from highsync_wrapper.runner import run_highsync
|
||||
|
||||
return run_highsync(
|
||||
highsync_model=highsync_model,
|
||||
images=images,
|
||||
audio=audio,
|
||||
input_fps=input_fps,
|
||||
max_frames=max_frames,
|
||||
inference_steps=None,
|
||||
cfg_scale=None,
|
||||
seed=42,
|
||||
preprocess_chunk_frames=1500,
|
||||
denoise_audio=False,
|
||||
trim_audio_to_video=True,
|
||||
keep_intermediate_outputs=True,
|
||||
)
|
||||
|
||||
|
||||
class HighSyncLipSyncAdvanced(HighSyncLipSync):
|
||||
@classmethod
|
||||
def INPUT_TYPES(cls):
|
||||
return {
|
||||
"required": {
|
||||
"highsync_model": ("HIGHSYNC_MODEL",),
|
||||
"images": ("IMAGE",),
|
||||
"audio": ("AUDIO",),
|
||||
"input_fps": ("INT", {"default": 25, "min": 1, "max": 120, "step": 1}),
|
||||
"max_frames": ("INT", {"default": 0, "min": 0, "max": 100000, "step": 1}),
|
||||
"inference_steps": ("INT", {"default": 20, "min": 1, "max": 100, "step": 1}),
|
||||
"cfg_scale": ("FLOAT", {"default": 3.5, "min": 0.0, "max": 20.0, "step": 0.1}),
|
||||
"seed": ("INT", {"default": 42, "min": 0, "max": 0xFFFFFFFF, "step": 1}),
|
||||
"preprocess_chunk_frames": ("INT", {"default": 1500, "min": 12, "max": 10000, "step": 12}),
|
||||
"denoise_audio": ("BOOLEAN", {"default": False}),
|
||||
"trim_audio_to_video": ("BOOLEAN", {"default": True}),
|
||||
"keep_intermediate_outputs": ("BOOLEAN", {"default": True}),
|
||||
}
|
||||
}
|
||||
|
||||
FUNCTION = "run_advanced"
|
||||
CATEGORY = "video/lipsync"
|
||||
|
||||
def run_advanced(
|
||||
self,
|
||||
highsync_model,
|
||||
images,
|
||||
audio,
|
||||
input_fps,
|
||||
max_frames,
|
||||
inference_steps,
|
||||
cfg_scale,
|
||||
seed,
|
||||
preprocess_chunk_frames,
|
||||
denoise_audio,
|
||||
trim_audio_to_video,
|
||||
keep_intermediate_outputs,
|
||||
):
|
||||
try:
|
||||
from .highsync_wrapper.runner import run_highsync
|
||||
except ImportError:
|
||||
from highsync_wrapper.runner import run_highsync
|
||||
|
||||
return run_highsync(
|
||||
highsync_model=highsync_model,
|
||||
images=images,
|
||||
audio=audio,
|
||||
input_fps=input_fps,
|
||||
max_frames=max_frames,
|
||||
inference_steps=inference_steps,
|
||||
cfg_scale=cfg_scale,
|
||||
seed=seed,
|
||||
preprocess_chunk_frames=preprocess_chunk_frames,
|
||||
denoise_audio=denoise_audio,
|
||||
trim_audio_to_video=trim_audio_to_video,
|
||||
keep_intermediate_outputs=keep_intermediate_outputs,
|
||||
)
|
||||
|
||||
|
||||
NODE_CLASS_MAPPINGS = {
|
||||
"HighSyncDownloadModels": HighSyncDownloadModels,
|
||||
"HighSyncModelLoader": HighSyncModelLoader,
|
||||
"HighSyncLipSync": HighSyncLipSync,
|
||||
"HighSyncLipSyncAdvanced": HighSyncLipSyncAdvanced,
|
||||
}
|
||||
|
||||
NODE_DISPLAY_NAME_MAPPINGS = {
|
||||
"HighSyncDownloadModels": "HighSync Download Models",
|
||||
"HighSyncModelLoader": "HighSync Model Loader",
|
||||
"HighSyncLipSync": "HighSync LipSync",
|
||||
"HighSyncLipSyncAdvanced": "HighSync LipSync Advanced",
|
||||
}
|
||||
@@ -0,0 +1,17 @@
|
||||
accelerate>=0.28.0
|
||||
diffusers>=0.24.0
|
||||
einops>=0.8.0
|
||||
huggingface-hub>=0.25.2
|
||||
omegaconf>=2.3.0
|
||||
safetensors
|
||||
transformers>=4.39.2
|
||||
torchvision
|
||||
torchaudio
|
||||
librosa>=0.10.2
|
||||
soundfile>=0.13.1
|
||||
opencv-python>=4.9.0.80
|
||||
pillow>=10.3.0
|
||||
imageio
|
||||
insightface>=0.7.3
|
||||
mediapipe>=0.10.14
|
||||
onnxruntime-gpu>=1.18.0
|
||||
@@ -0,0 +1,36 @@
|
||||
from pathlib import Path
|
||||
|
||||
import pytest
|
||||
|
||||
from highsync_wrapper.download_utils import validate_required_models
|
||||
|
||||
|
||||
def test_validate_required_models_reports_missing_files(tmp_path):
|
||||
with pytest.raises(FileNotFoundError) as exc_info:
|
||||
validate_required_models(tmp_path)
|
||||
|
||||
message = str(exc_info.value)
|
||||
assert "denoising_unet-500.pth" in message
|
||||
assert "reference_unet-500.pth" in message
|
||||
assert "audio_processor/whisper_tiny.pt" in message
|
||||
assert "sd-vae-ft-mse/config.json" in message
|
||||
assert "sd-image-variations-diffusers/unet/config.json" in message
|
||||
|
||||
|
||||
def test_validate_required_models_accepts_expected_layout(tmp_path):
|
||||
(tmp_path / "denoising_unet-500.pth").write_bytes(b"")
|
||||
(tmp_path / "reference_unet-500.pth").write_bytes(b"")
|
||||
|
||||
audio_dir = tmp_path / "audio_processor"
|
||||
audio_dir.mkdir()
|
||||
(audio_dir / "whisper_tiny.pt").write_bytes(b"")
|
||||
|
||||
vae_dir = tmp_path / "sd-vae-ft-mse"
|
||||
vae_dir.mkdir()
|
||||
(vae_dir / "config.json").write_text("{}", encoding="utf-8")
|
||||
|
||||
unet_dir = tmp_path / "sd-image-variations-diffusers" / "unet"
|
||||
unet_dir.mkdir(parents=True)
|
||||
(unet_dir / "config.json").write_text("{}", encoding="utf-8")
|
||||
|
||||
assert validate_required_models(tmp_path) == Path(tmp_path).resolve()
|
||||
Reference in New Issue
Block a user