mirror of
https://github.com/THUDM/CogVideo.git
synced 2026-05-08 07:58:10 +08:00
207020afea
Add --device_map CLI argument to enable multi-GPU inference without
code modifications. This implements the functionality described in
existing comments at lines 92-93 and 146-150.
Changes:
- Add device_map parameter to generate_video() with options:
- None (default): Uses sequential CPU offload (backward compatible)
- 'balanced': Distributes model evenly across GPUs (recommended)
- 'auto': Automatic device placement by accelerate
- 'sequential': Fills GPUs one by one in order
- Conditionally enable CPU offload only when device_map is None
- Pass device_map to from_pretrained() for all three pipeline types:
CogVideoXPipeline, CogVideoXImageToVideoPipeline, CogVideoXVideoToVideoPipeline
- Update module docstring with multi-GPU usage examples
- Add validation for device_map values
Backward compatibility:
- Default behavior unchanged (device_map=None uses CPU offload)
- All existing scripts work without modification
Usage:
# Multi-GPU with balanced distribution (recommended):
python cli_demo.py --prompt '...' --model_path THUDM/CogVideoX1.5-5b \
--device_map balanced
# Multi-GPU with auto placement:
python cli_demo.py --prompt '...' --model_path THUDM/CogVideoX1.5-5b \
--device_map auto