Skip to content

feat: Add support for Blackwell GPU architecture (sm_120) and fix compatibility issues#416

Open
Sunmingyang1987 wants to merge 2 commits into
TMElyralab:mainfrom
Sunmingyang1987:blackwell-support
Open

feat: Add support for Blackwell GPU architecture (sm_120) and fix compatibility issues#416
Sunmingyang1987 wants to merge 2 commits into
TMElyralab:mainfrom
Sunmingyang1987:blackwell-support

Conversation

@Sunmingyang1987

Copy link
Copy Markdown

Summary

This PR adds support for NVIDIA Blackwell GPU architecture (sm_120), such as RTX PRO 6000 Blackwell Max-Q, and fixes several compatibility issues when running on newer PyTorch/CUDA environments.

Changes

Code Changes

  • musetalk/utils/preprocessing.py: Add monkey-patch for torch.load to handle PyTorch 2.6+ weights_only=True change, allowing legacy checkpoint (dwpose) loading
  • musetalk/utils/face_detection/detection/sfd/sfd_detector.py: Move SFD model path to models/s3fd/ for unified model management
  • scripts/inference.py: Fix --ffmpeg_path parameter being ignored when system ffmpeg is already available in PATH

Documentation (README.md)

  • Added SFD model (models/s3fd/s3fd.pth) to model directory structure
  • Added preprocessor_config.json note for whisper model
  • Added comprehensive Troubleshooting section covering:
    • Blackwell GPU (sm_120) support: PyTorch upgrade + mmcv 2.1.0 installation
    • PyTorch 2.6+ checkpoint loading error
    • FFmpeg missing PNG encoder

Tested Environment

  • GPU: NVIDIA RTX PRO 6000 Blackwell Max-Q Workstation Edition (sm_120)
  • PyTorch: 2.9.0 + CUDA 12.8
  • mmcv: 2.1.0 (compiled from source)
  • mmdet: 3.1.0
  • mmpose: 1.1.0

Sign up for free to join this conversation on GitHub. Already have an account? Sign in to comment

Labels

None yet

Projects

None yet

Development

Successfully merging this pull request may close these issues.

1 participant