DOCSEARCHbyAlgolia
Submit Docs

StableAvatar Documentation

repository·main·Indexed 23 days ago

https://github.com/francis-rings/stableavatar

StableAvatar is an end-to-end video diffusion transformer for generating infinite-length, high-fidelity, audio-driven avatar videos. It preserves identity without post-processing tools like FaceFusion or GFP-GAN. Key features include a Time-step-aware Audio Adapter, an Audio Native Guidance Mechanism, and a Dynamic Weighted Sliding-window Strategy. The system supports various resolutions (512x512, 480x832, 832x480), multiple GPU memory modes for VRAM optimization, and options for full or LoRA fine-tuning based on Wan2.1 and Wav2Vec2.0.

Tokens
2.6K
Snippets
5
Records
14
Agent score
29%

What's inside StableAvatar

  1. Overview of StableAvatar

    main

    StableAvatar is an end-to-end video diffusion transformer designed for infinite-length, high-quality, audio-driven avatar video generation. Unlike existing models that require post-processing (like face-swapping or face restoration), StableAvatar synthesizes identity-preserving videos directly.

    Key technical features include:

    • Time-step-aware Audio Adapter: Prevents latent distribution error accumulation during long video generation.
    • Audio Native Guidance Mechanism: Enhances audio synchronization by using the diffusion model's evolving joint audio-latent prediction as a dynamic guidance signal.
    • Dynamic Weighted Sliding-window Strategy: Fuses latents over time to ensure smoothness in infinite-length videos.

  2. Prepare training dataset for StableAvatar

    main

    The training dataset must follow a specific directory structure organized by resolution (square, rec, or vec) and activity type (speech, singing, dancing).

    Directory Structure:

    talking_face_data/
├── [square|rec|vec]/
│   ├── [speech|singing|dancing]/
│   │   └── 00001/
│   │       ├── sub_clip.mp4
│   │       ├── audio.wav
│   │       ├── images/ (contains frame_i.png)
│   │       ├── face_masks/ (contains frame_i.png)
│   │       └── lip_masks/ (contains frame_i.png)

    Data Extraction Tasks:

    • Frames: Use ffmpeg to extract frames from raw videos into the images folder.
    • Face Masks: Use the StableAnimator repository.
    • Lip Masks: Use lip_mask_extractor.py with mediapipe.

    Lip Mask Extraction Command:

    pip install mediapipe
python lip_mask_extractor.py --folder_root="path/to/dataset" --start=1 --end=500

    Path Index Files:

    Maintain text files (video_square_path.txt, video_rec_path.txt, video_vec_path.txt) that list the full paths to each video folder (e.g., path/to/talking_face_data/rec/speech/00001).

  3. Run Base Model inference

    main

    Run inference using the provided inference.sh script. You can also launch a Gradio web interface using python app.py.

    Key Configuration Options in inference.sh:

    • --width / --height: Set resolution (supports 512x512, 480x832, or 832x480).
    • --output_dir: Path to save generated animations.
    • --validation_reference_path: Path to the reference image.
    • --validation_driven_audio_path: Path to the driving audio.
    • --validation_prompts: Text prompts. Recommended format: [Description of first frame]-[Description of human behavior]-[Description of background (optional)].
    • --pretrained_model_name_or_path: Path to pretrained Wan2.1-1.3B weights.
    • --pretrained_wav2vec_path: Path to pretrained Wav2Vec2.0 weights.
    • --transformer_path: Path to pretrained StableAvatar weights (can switch between transformer3d-square.pt and transformer3d-rec-vec.pt).
    • --sample_steps: Total inference steps. Recommended: [30-50].
    • --overlap_window_length: Overlapping context length. Recommended: [5-15].
    • --clip_sample_n_frames: Synthesized frames per batch/context window.
    • --sample_text_guide_scale / --sample_audio_guide_scale: CFG scales. Recommended: [3-6]. Increasing audio CFG helps lip sync.

    Memory Management (--GPU_memory_mode):

    • model_full_load: Standard mode.
    • sequential_cpu_offload: Uses ~3GB VRAM, but slower.
    • model_cpu_offload_and_qfloat8: Optimized memory.
    • model_cpu_offload: Reduces VRAM usage by roughly half compared to model_full_load.

    Multi-GPU Inference:

    To speed up inference on multiple GPUs, modify --ulysses_degree and --ring_degree. Ensure ulysses_degree * ring_degree = total GPU number / world-size. Example for 8 GPUs: --ulysses_degree=4 --ring_degree=2. Use bash multiple_gpu_inference.sh for pre-configured multi-GPU setups.

  4. Install StableAvatar environment

    main

    To set up the environment for the basic Wan2.1-1.3B-based model, install the required PyTorch versions and dependencies. If you want to accelerate attention computation, you can optionally install flash_attn.

    Standard Setup:

    pip install torch==2.6.0 torchvision==0.21.0 torchaudio==2.1.1 --index-url https://download.pytorch.org/whl/cu124
pip install -r requirements.txt
# Optional
pip install flash_attn

    Blackwell Series Chips Setup:

    pip install torch==2.7.0 torchvision==0.22.0 torchaudio==2.7.0 --index-url https://download.pytorch.org/whl/cu128
pip install -r requirements.txt
# Optional
pip install flash_attn
    pip install torch==2.6.0 torchvision==0.21.0 torchaudio==2.1.1 --index-url https://download.pytorch.org/whl/cu124
pip install -r requirements.txt
# Optional to install flash_attn to accelerate attention computation
pip install flash_attn

  5. Download StableAvatar weights

    main

    Download the model weights using the Hugging Face CLI. If you are in a region with connection issues to Hugging Face, set the HF_ENDPOINT environment variable to use a mirror.

    1. Install the CLI: pip install "huggingface_hub[cli]"
    2. Set mirror (optional): export HF_ENDPOINT=https://hf-mirror.com
    3. Download to the checkpoints directory.
    pip install "huggingface_hub[cli]"
cd StableAvatar
mkdir checkpoints
huggingface-cli download FrancisRing/StableAvatar --local-dir ./checkpoints

  6. Train StableAvatar

    main

    Training can be performed on single or multiple machines for different resolution settings.

    Training Modes:

    • Single Resolution (512x512): Use bash train_1B_square.sh.
    • Mixed Resolution (480x832 and 832x480): Use bash train_1B_rec_vec.sh.
    • Multi-Machine Training: Use scripts ending in _64.sh (e.g., train_1B_square_64.sh).
    • LoRA Training: Use scripts like train_1B_rec_vec_lora.sh. Control quality via --rank and --network_alpha.
    • Wan2.1-14B Training: Use bash train_14B.sh (requires downloading 14B weights first).

    Key Training Parameters:

    • --pretrained_model_name_or_path: Path to Wan2.1-1.3B weights.
    • --pretrained_wav2vec_path: Path to Wav2Vec2.0 weights.
    • --output_dir: Path where checkpoints are saved.
    • --train_data_square_dir / --train_data_rec_dir / --train_data_vec_dir: Paths to the .txt files containing dataset folder paths.
    • --video_sample_n_frames: Number of frames processed in a single batch.
    • --num_train_epochs: Number of epochs (default is infinite; terminate manually).

    VRAM Requirements:

    • Mixed Resolution: ~50GB VRAM.
    • 512x512 Only: ~40GB VRAM.

  7. Fine-tune StableAvatar

    main

    You can perform full fine-tuning or LoRA fine-tuning by providing a --transformer_path pointing to an existing StableAvatar checkpoint.

    Full Fine-tuning: Add --transformer_path="path/to/checkpoint/transformer3d-square.pt" to your training script (e.g., train_1B_rec_vec.sh).

    LoRA Fine-tuning: Add the same --transformer_path to the LoRA training script (e.g., train_1B_rec_vec_lora.sh). Use --rank and --network_alpha to tune quality.

  11. Separate vocals from audio

    main

    To improve lip synchronization, you can separate the vocal signal from a noisy audio file using vocal_seperator.py. This requires the audio-separator[gpu] package and the Kim_Vocal_2.onnx model checkpoint.

    Requirements:

    pip install audio-separator[gpu]

    Usage:

    python vocal_seperator.py --audio_separator_model_file="path/StableAvatar/checkpoints/Kim_Vocal_2.onnx" --audio_file_path="path/test/audio.wav" --saved_vocal_path="path/test/vocal.wav"
    pip install audio-separator[gpu]
python vocal_seperator.py --audio_separator_model_file="path/StableAvatar/checkpoints/Kim_Vocal_2.onnx" --audio_file_path="path/test/audio.wav" --saved_vocal_path="path/test/vocal.wav"

  12. Configure GPU Memory Modes

    main

    The generate function supports different memory management strategies via the GPU_memory_mode parameter to accommodate different VRAM capacities:

    • Normal: Uses approximately 25GB VRAM. All components are moved to the GPU.
    • model_cpu_offload: Uses approximately 13GB VRAM. Enables model CPU offloading.
    • model_cpu_offload_and_qfloat8: Uses model CPU offloading combined with FP8 weight conversion for the transformer to reduce memory footprint.
    • sequential_cpu_offload: Uses sequential CPU offloading for even lower memory requirements.