MusePose Documentation

repository·main·Indexed 25 days ago

https://github.com/tmelyralab/musepose

A pose-driven image-to-video generation framework for virtual human generation. MusePose enables the animation of a character from a reference image using a provided pose sequence. The framework includes tools for pose alignment, inference with VRAM optimization, and a multi-stage training process utilizing MMLab packages and accelerate.

Tokens
1.1K
Snippets
3
Records
5
Agent score
33%

What's inside MusePose

  1. Organize MusePose Pretrained Weights

    main

    MusePose requires several pretrained weights to be organized in a specific directory structure under ./pretrained_weights/.

    Required Weights:

    • MusePose: denoising_unet.pth, motion_module.pth, pose_guider.pth, reference_unet.pth (from HuggingFace).
    • dwpose: dw-ll_ucoco_384.pth and yolox_l_8x8_300e_coco.pth (Note: rename the downloaded yolox file to yolox_l_8x8_300e_coco.pth).
    • sd-image-variations-diffusers: unet/config.json and unet/diffusion_pytorch_model.bin.
    • image_encoder: config.json and pytorch_model.bin.
    • sd-vae-ft-mse: config.json and diffusion_pytorch_model.bin.

    Note: control_v11p_sd15_openpose and animatediff weights are required for training only.

    ./pretrained_weights/
    |-- MusePose
    |   |-- denoising_unet.pth
    |   |-- motion_module.pth
    |   |-- pose_guider.pth
    |   └── reference_unet.pth
    |-- dwpose
    |   |-- dw-ll_ucoco_384.pth
    |   |-- yolox_l_8x8_300e_coco.pth
    |-- sd-image-variations-diffusers
    |   └── unet
    |       |-- config.json
    |       └── diffusion_pytorch_model.bin
    |-- image_encoder
    |   |-- config.json
    |   └── pytorch_model.bin
    |-- sd-vae-ft-mse
    |   |-- config.json
    |   └── diffusion_pytorch_model.bin
    |-- control_v11p_sd15_openpose
    |   └── diffusion_pytorch_model.bin
    └── animatediff
        └── mm_sd_v15_v2.ckpt
  2. Install MusePose and Build Environment

    main

    To install MusePose, ensure you are using Python >=3.10 and CUDA 11.7. First, install the base requirements using pip.

    Then, install the necessary MMLab packages using openmim to ensure compatibility with the model's dependencies.

    # Install base requirements
    pip install -r requirements.txt
    
    # Install MMLab packages
    pip install --no-cache-dir -U openmim 
    mim install mmengine 
    mim install "mmcv>=2.0.1" 
    mim install "mmdet>=3.1.0" 
    mim install "mmpose>=1.1.0" 
  3. Train MusePose Models

    main

    Training involves a multi-step data preparation and execution process.

    1. Data Preparation

    1. Extract keypoints: python extract_dwpose_keypoints.py --video_dir ./your_videos_dir
    2. Render dwpose: python draw_dwpose.py --video_dir ./your_videos_dir
    3. Generate metadata: python extract_meta_info_multiple_dataset.py --video_dirs ./your_videos_dir --dataset_name your_dataset_name

    2. Environment Setup

    Install accelerate and configure your machine (e.g., for DeepSpeed) using:

    pip install accelerate
    accelerate config

    3. Launch Training

    Use accelerate launch with the appropriate stage configuration file.

    • Stage 1: ./configs/train_stage_1.yaml
    • Stage 2: ./configs/train_stage_2.yaml
  4. Run MusePose Inference

    main

    To run inference, follow these steps:

    1. Configure the test file: Edit ./configs/test_stage_2.yaml to map your reference image to the aligned pose video path:
      test_cases:
        "./assets/images/ref.png":
          - "./assets/poses/align/img_ref_video_dance.mp4"
    2. Execute inference: Run the test_stage_2.py script.

    VRAM Optimization: You can reduce VRAM usage by specifying a lower resolution using -W (width) and -H (height). The output will be resized back to the original pose video dimensions.

    • 512x512 requires ~16GB VRAM.
    • 768x768 requires ~28GB VRAM.
  5. Perform Pose Alignment for Inference

    main

    Before running inference, you must align the pose of the reference image to the target dance video using the pose_align.py script. This generates the necessary pose sequence for the model.

    Input files should be organized in ./assets/images/ and ./assets/videos/.

    python pose_align.py --imgfn_refer ./assets/images/ref.png --vidfn ./assets/videos/dance.mp4