repository·main·Indexed 19 days ago
https://github.com/sstzal/difftalkA PyTorch implementation of a CVPR 2023 paper utilizing Diffusion Models for generalized audio-driven portrait animation. DiffTalk focuses on high-quality talking head generation from audio signals, employing an LDM-based iterative denoising process. The project includes tools for training and inference via main.py, supporting HDTF dataset preprocessing, OmegaConf-based YAML configurations, and linear learning rate scaling.
DiffTalk requires Python 3.7.0 and specific versions of PyTorch and PyTorch Lightning. Experiments are conducted using 8 NVIDIA 3090Ti GPUs. Ensure you have the following dependencies installed:
Refer to requirements.txt for the full dependency list.
# Ensure these versions are met:
python 3.7.0
pytorch 1.10.0
pytorch-lightning 1.2.5
torchvision 0.11.0To use a pretrained model for the first stage of the DiffTalk pipeline, you must place the model checkpoint file in the models/ directory and ensure it is named exactly model.ckpt.
mv path/to/your/pretrained_model.ckpt models/model.ckptTo use the HDTF dataset, you must preprocess it following these steps:
./data/HDTF using the following structure:./data/HDTF
|——images
| |——0_0.jpg
| |——N_M.bin
|——landmarks
| |——0_0.lmd
| |——N_M.lms
|——audio_smooth
| |——0_0.npy
| |——N_M.npydata_train.txt and data_test.txt files containing the identifiers (e.g., 0_0, 0_1, ..., N_M) where N is the total number of classes and M is the class size.Before running the project, download the first stage model and place it in the ./models directory.
# Place the downloaded model in:
./modelsThe project implements a linear scaling rule for the learning rate to maintain stability when changing batch sizes or GPU counts.
If --scale_lr is enabled (default), the model's learning rate is calculated as:
model.learning_rate = accumulate_grad_batches * ngpu * batch_size * base_learning_rate
Where:
accumulate_grad_batches: Number of gradient accumulation steps.ngpu: Number of GPUs being used.batch_size: The per-device batch size.base_learning_rate: The value provided in the model.base_learning_rate config.The main.py script serves as the primary entrypoint for training and testing models. It uses a combination of YAML configuration files and command-line arguments to define the model, data, and trainer settings.
Configuration Hierarchy:
-b/--base, merged left-to-right).--key value or nested.key=value syntax).Key CLI Flags:
-t, --train: Enables training mode.--no-test: Disables the testing phase after training.-r, --resume: Resumes training from a specific log directory or checkpoint.-r2, --resume2: Loads a state dict from a specific checkpoint file into the model (non-strict).-n, --name: Adds a postfix to the log directory name.-l, --logdir: Specifies the base directory for logging (default: logs).-s, --seed: Sets the random seed (default: 23).-d, --debug: Enables post-mortem debugging on failure.# Example: Train using a base config and override a parameter
python main.py -b configs/base_config.yaml --model.params.lr=0.0001 -t
# Example: Resume training from a log directory
python main.py -r logs/2023-01-01T12-00-00_config_name -tUsers should be aware of the following limitations:
z_T must completely cover the mouth region and must not leak any lip shape information.To start the training process, execute the provided shell script.
sh run.shTo test the model and perform inference, execute the inference shell script.
sh inference.shThe project uses OmegaConf to manage complex configurations. The expected YAML structure for a complete configuration includes the following top-level keys:
model: Defines the model architecture and hyperparameters.base_learning_rate: The base LR used for scaling.target: The import path to the Lightning module.params: Dictionary of parameters for the model.data: Defines the DataModuleFromConfig settings.target: Usually main.DataModuleFromConfig.params: Contains batch_size, train, validation, test, etc.lightning: (Optional) Configuration for the PyTorch Lightning Trainer and its components.trainer: Arguments passed to the Trainer (e.g., gpus, accumulate_grad_batches).logger: Configuration for the logger (e.g., WandbLogger or TestTubeLogger).modelcheckpoint: Configuration for ModelCheckpoint.callbacks: A dictionary of callbacks to instantiate.model:
base_learning_rate: 0.0001
target: ldm.models.diffusion.ddpm.LatentDiffusion
params:
some_param: value
data:
target: main.DataModuleFromConfig
params:
batch_size: 4
train:
target: my_dataset_module
params:
path: /data/train
lightning:
trainer:
gpus: 2
accumulate_grad_batches: 4
logger:
target: pytorch_lightning.loggers.WandbLogger
params:
name: my_experimentThe following arguments are available via the command line:
| Flag | Long Flag | Type | Description |
|---|---|---|---|
-n | --name | str | Postfix for logdir |
-r | --resume | str | Resume from logdir or checkpoint in logdir |
-r2 | --resume2 | str | Resume from logdir or checkpoint in logdir |
-b | --base | list | Paths to base configs (loaded left-to-right) |
-t | --train | bool | Enable training |
--no-test | --no-test | bool | Disable test phase |
-p | --project | str | Name of new or path to existing project |
-d | --debug | bool | Enable post-mortem debugging |
-s | --seed | int | Seed for seed_everything (default: 23) |
-f | --postfix | str | Post-postfix for default name |
-l | --logdir | str | Directory for logging (default: logs) |
--scale_lr | --scale_lr | bool | Scale base-lr by ngpu * batch_size * n_accumulate (default: True) |
The ImageLogger callback is used to log images during training and validation. It can log to local files or to a logger like TestTubeLogger.
Key Parameters:
batch_frequency: How often to log (in batches).max_images: Maximum number of images to log per batch.clamp: If True, clamps images to [-1, 1].rescale: If True, rescales images from [-1, 1] to [0, 1] for visualization.log_on_batch_idx: If True, uses batch_idx for frequency check; otherwise uses global_step.log_first_step: If True, logs the very first step regardless of frequency.