Zengtudor/Wan2.1
1
0
Fork 0
mirror of https://github.com/Wan-Video/Wan2.1.git synced 2025-12-15 11:43:21 +00:00
Code Issues Actions Packages Projects Releases Wiki Activity

Merge 28a931100b into 854bd88e7f

Browse Source
This commit is contained in:
Kyle R 2025-12-15 17:35:28 +08:00 committed by GitHub
commit 6f4387fd34
No known key found for this signature in database
GPG Key ID: B5690EEEBB952194
24 changed files with 3171 additions and 3 deletions
Show Stats Download Patch File Download Diff File Expand all files Collapse all files

163
.github/ISSUE_TEMPLATE/bug_report.yml vendored Normal file
View File

@ -0,0 +1,163 @@
name: Bug Report
description: File a bug report to help us improve
title: "[Bug]: "
labels: ["bug", "needs-triage"]
body:
- type: markdown
attributes:
value: |
Thanks for taking the time to fill out this bug report! Please provide as much detail as possible.
- type: textarea
id: description
attributes:
label: Bug Description
description: A clear and concise description of the bug
placeholder: What went wrong?
validations:
required: true
- type: textarea
id: reproduce
attributes:
label: Steps to Reproduce
description: Steps to reproduce the behavior
placeholder: |
1. Load model with '...'
2. Run inference with '...'
3. See error
validations:
required: true
- type: textarea
id: expected
attributes:
label: Expected Behavior
description: What you expected to happen
placeholder: What should have happened?
validations:
required: true
- type: textarea
id: actual
attributes:
label: Actual Behavior
description: What actually happened
placeholder: What actually happened?
validations:
required: true
- type: textarea
id: logs
attributes:
label: Error Logs
description: Please copy and paste any relevant error messages or logs
render: shell
- type: dropdown
id: pipeline
attributes:
label: Pipeline
description: Which pipeline are you using?
options:
- Text-to-Video (T2V)
- Image-to-Video (I2V)
- First-Last-Frame-to-Video (FLF2V)
- VACE (Video Creation & Editing)
- Text-to-Image (T2I)
- Other
validations:
required: true
- type: input
id: version
attributes:
label: Wan2.1 Version
description: What version of Wan2.1 are you using?
placeholder: "2.1.0"
validations:
required: true
- type: dropdown
id: model-size
attributes:
label: Model Size
description: Which model size are you using?
options:
- 14B
- 1.3B
- Not applicable
validations:
required: true
- type: input
id: python-version
attributes:
label: Python Version
description: What version of Python are you using?
placeholder: "3.10.0"
validations:
required: true
- type: input
id: pytorch-version
attributes:
label: PyTorch Version
description: What version of PyTorch are you using?
placeholder: "2.4.0"
validations:
required: true
- type: input
id: cuda-version
attributes:
label: CUDA Version
description: What version of CUDA are you using? (or N/A for CPU)
placeholder: "11.8"
- type: dropdown
id: gpu
attributes:
label: GPU Type
description: What GPU are you using?
options:
- NVIDIA A100
- NVIDIA V100
- NVIDIA RTX 4090
- NVIDIA RTX 3090
- NVIDIA RTX 3080
- Other NVIDIA GPU
- AMD GPU
- CPU only
- Other
- type: textarea
id: environment
attributes:
label: Environment Details
description: Any additional environment details
placeholder: |
- OS: Ubuntu 22.04
- RAM: 64GB
- Number of GPUs: 2
- Other relevant details
- type: textarea
id: additional
attributes:
label: Additional Context
description: Add any other context about the problem here
placeholder: Screenshots, videos, or additional information
- type: checkboxes
id: checklist
attributes:
label: Checklist
description: Please confirm the following
options:
- label: I have searched existing issues to ensure this is not a duplicate
required: true
- label: I have provided all required information
required: true
- label: I have included error logs (if applicable)
required: false

118
.github/ISSUE_TEMPLATE/feature_request.yml vendored Normal file
View File

@ -0,0 +1,118 @@
name: Feature Request
description: Suggest a new feature or enhancement
title: "[Feature]: "
labels: ["enhancement", "needs-triage"]
body:
- type: markdown
attributes:
value: |
Thanks for suggesting a feature! Please provide as much detail as possible to help us understand your request.
- type: textarea
id: problem
attributes:
label: Problem Statement
description: Is your feature request related to a problem? Please describe.
placeholder: I'm frustrated when...
validations:
required: true
- type: textarea
id: solution
attributes:
label: Proposed Solution
description: Describe the solution you'd like
placeholder: I would like to see...
validations:
required: true
- type: textarea
id: alternatives
attributes:
label: Alternatives Considered
description: Describe any alternative solutions or features you've considered
placeholder: I also considered...
- type: dropdown
id: feature-type
attributes:
label: Feature Type
description: What type of feature is this?
options:
- New Pipeline/Model
- Performance Improvement
- API Enhancement
- Documentation
- Developer Experience
- Infrastructure
- Other
validations:
required: true
- type: dropdown
id: priority
attributes:
label: Priority
description: How important is this feature to you?
options:
- Critical - Blocking my work
- High - Needed soon
- Medium - Would be nice to have
- Low - Nice to have eventually
validations:
required: true
- type: textarea
id: use-case
attributes:
label: Use Case
description: Describe your use case for this feature
placeholder: |
I want to use this feature to...
This would help me...
validations:
required: true
- type: textarea
id: implementation
attributes:
label: Implementation Ideas
description: If you have ideas about how to implement this, please share
placeholder: |
This could be implemented by...
Potential challenges might include...
- type: textarea
id: examples
attributes:
label: Examples
description: Provide code examples or mockups of how this feature would work
render: python
- type: checkboxes
id: contribution
attributes:
label: Contribution
description: Would you be willing to contribute to this feature?
options:
- label: I would like to implement this feature
- label: I can help test this feature
- label: I can help with documentation
- type: textarea
id: additional
attributes:
label: Additional Context
description: Add any other context, screenshots, or examples
placeholder: Links to similar features in other projects, mockups, etc.
- type: checkboxes
id: checklist
attributes:
label: Checklist
description: Please confirm the following
options:
- label: I have searched existing issues to ensure this is not a duplicate
required: true
- label: I have clearly described the problem and proposed solution
required: true

94
.github/dependabot.yml vendored Normal file
View File

@ -0,0 +1,94 @@
# Dependabot configuration for automated dependency updates
# Documentation: https://docs.github.com/en/code-security/dependabot
version: 2
updates:
# Python dependencies
- package-ecosystem: "pip"
directory: "/"
schedule:
interval: "weekly"
day: "monday"
time: "09:00"
open-pull-requests-limit: 10
reviewers:
- "kuaishou/wan-maintainers" # Update with actual team
assignees:
- "kuaishou/wan-maintainers" # Update with actual team
labels:
- "dependencies"
- "python"
commit-message:
prefix: "deps"
prefix-development: "deps-dev"
include: "scope"
# Group minor and patch updates together
groups:
pytorch-ecosystem:
patterns:
- "torch*"
- "torchvision"
update-types:
- "minor"
- "patch"
transformers-ecosystem:
patterns:
- "transformers"
- "diffusers"
- "accelerate"
update-types:
- "minor"
- "patch"
dev-dependencies:
dependency-type: "development"
update-types:
- "minor"
- "patch"
# Ignore specific dependencies that need manual updates
ignore:
# Flash attention requires specific CUDA versions
- dependency-name: "flash-attn"
update-types: ["version-update:semver-major"]
# PyTorch major updates require testing
- dependency-name: "torch"
update-types: ["version-update:semver-major"]
- dependency-name: "torchvision"
update-types: ["version-update:semver-major"]
# Allow specific versions
allow:
- dependency-type: "direct"
- dependency-type: "production"
- dependency-type: "development"
# GitHub Actions
- package-ecosystem: "github-actions"
directory: "/"
schedule:
interval: "weekly"
day: "monday"
time: "09:00"
open-pull-requests-limit: 5
reviewers:
- "kuaishou/wan-maintainers"
labels:
- "dependencies"
- "github-actions"
commit-message:
prefix: "ci"
include: "scope"
groups:
github-actions:
patterns:
- "*"
update-types:
- "minor"
- "patch"
# Docker (if Dockerfile exists)
# - package-ecosystem: "docker"
# directory: "/"
# schedule:
# interval: "weekly"
# labels:
# - "dependencies"
# - "docker"

128
.github/pull_request_template.md vendored Normal file
View File

@ -0,0 +1,128 @@
## Description
<!-- Provide a brief description of the changes in this PR -->
## Type of Change
<!-- Mark the relevant option with an 'x' -->
- [ ] Bug fix (non-breaking change which fixes an issue)
- [ ] New feature (non-breaking change which adds functionality)
- [ ] Breaking change (fix or feature that would cause existing functionality to not work as expected)
- [ ] Documentation update
- [ ] Performance improvement
- [ ] Code refactoring
- [ ] Test addition/modification
- [ ] CI/CD changes
- [ ] Dependency update
## Related Issues
<!-- Link to related issues using #issue_number -->
Closes #
Relates to #
## Changes Made
<!-- Provide a detailed list of changes -->
-
-
-
## Testing
### Test Environment
- Python version:
- PyTorch version:
- CUDA version:
- GPU type:
- Number of GPUs:
### Testing Performed
<!-- Describe the tests you ran and their results -->
- [ ] All existing tests pass
- [ ] Added new unit tests
- [ ] Added new integration tests
- [ ] Manual testing completed
- [ ] Tested on CPU
- [ ] Tested on GPU
- [ ] Tested with 14B model
- [ ] Tested with 1.3B model
### Test Results
<!-- Paste relevant test output -->
```
pytest output here
```
## Performance Impact
<!-- If applicable, describe any performance changes -->
- Inference speed:
- Memory usage:
- GPU utilization:
## Breaking Changes
<!-- If this is a breaking change, describe what breaks and migration steps -->
-
-
## Documentation
<!-- Mark the relevant options with an 'x' -->
- [ ] README.md updated
- [ ] INSTALL.md updated
- [ ] Code comments added/updated
- [ ] Docstrings added/updated
- [ ] API documentation updated
- [ ] CHANGELOG.md updated
- [ ] No documentation needed
## Checklist
<!-- Ensure all items are completed before requesting review -->
- [ ] My code follows the project's style guidelines (YAPF/Black formatted)
- [ ] I have performed a self-review of my code
- [ ] I have commented my code, particularly in hard-to-understand areas
- [ ] I have made corresponding changes to the documentation
- [ ] My changes generate no new warnings
- [ ] I have added tests that prove my fix is effective or that my feature works
- [ ] New and existing unit tests pass locally with my changes
- [ ] Any dependent changes have been merged and published
- [ ] I have run `make format` to format the code
- [ ] I have checked my code with `mypy` for type errors
- [ ] I have updated type hints where necessary
- [ ] Pre-commit hooks pass
## Screenshots/Videos
<!-- If applicable, add screenshots or videos to demonstrate the changes -->
## Additional Notes
<!-- Add any additional notes, concerns, or context for reviewers -->
## Reviewer Notes
<!-- Anything specific you want reviewers to focus on? -->
---
**For Maintainers:**
- [ ] Code review completed
- [ ] Tests pass in CI
- [ ] Documentation is adequate
- [ ] Ready to merge

198
.github/workflows/ci.yml vendored Normal file
View File

@ -0,0 +1,198 @@
name: CI/CD Pipeline
on:
push:
branches: [ main, dev, 'claude/**' ]
pull_request:
branches: [ main, dev ]
jobs:
lint:
name: Code Quality & Linting
runs-on: ubuntu-latest
steps:
- uses: actions/checkout@v4
- name: Set up Python
uses: actions/setup-python@v5
with:
python-version: '3.10'
cache: 'pip'
- name: Install dependencies
run: |
python -m pip install --upgrade pip
pip install yapf black isort mypy
- name: Check formatting with YAPF
run: |
yapf --diff --recursive wan/ tests/
continue-on-error: true
- name: Check formatting with Black
run: |
black --check wan/ tests/
continue-on-error: true
- name: Check import sorting with isort
run: |
isort --check-only wan/ tests/
continue-on-error: true
- name: Type check with mypy
run: |
mypy wan/
continue-on-error: true
test-cpu:
name: CPU Tests
runs-on: ubuntu-latest
strategy:
matrix:
python-version: ['3.10', '3.11']
steps:
- uses: actions/checkout@v4
- name: Set up Python ${{ matrix.python-version }}
uses: actions/setup-python@v5
with:
python-version: ${{ matrix.python-version }}
cache: 'pip'
- name: Install system dependencies
run: |
sudo apt-get update
sudo apt-get install -y libavformat-dev libavcodec-dev libavutil-dev libswscale-dev
- name: Install Python dependencies (CPU-only)
run: |
python -m pip install --upgrade pip
pip install torch torchvision --index-url https://download.pytorch.org/whl/cpu
pip install -e .[dev]
- name: Run unit tests
run: |
pytest tests/ -v -m "not cuda and not requires_model and not integration" --tb=short
- name: Run import tests
run: |
python -c "from wan.modules.model import WanModel; print('WanModel import OK')"
python -c "from wan.modules.vae import WanVAE_; print('WanVAE import OK')"
python -c "from wan.modules.attention import attention; print('attention import OK')"
python -c "from wan.text2video import WanT2V; print('WanT2V import OK')"
python -c "from wan.image2video import WanI2V; print('WanI2V import OK')"
test-gpu:
name: GPU Tests (CUDA)
runs-on: ubuntu-latest
# Note: This requires a self-hosted runner with GPU access
# For public CI, this job can be skipped
if: false # Disable by default (enable for self-hosted runners with GPU)
steps:
- uses: actions/checkout@v4
- name: Set up Python
uses: actions/setup-python@v5
with:
python-version: '3.10'
- name: Install CUDA dependencies
run: |
# Add CUDA installation steps here
echo "CUDA installation required"
- name: Install Python dependencies
run: |
python -m pip install --upgrade pip
pip install torch torchvision --index-url https://download.pytorch.org/whl/cu118
pip install -e .[dev]
- name: Run GPU tests
run: |
pytest tests/ -v -m "cuda" --tb=short
security:
name: Security Scanning
runs-on: ubuntu-latest
steps:
- uses: actions/checkout@v4
- name: Set up Python
uses: actions/setup-python@v5
with:
python-version: '3.10'
- name: Install dependencies
run: |
python -m pip install --upgrade pip
pip install safety bandit
- name: Run safety check
run: |
pip install -e .
safety check --json || true
continue-on-error: true
- name: Run bandit security scan
run: |
bandit -r wan/ -f json || true
continue-on-error: true
build:
name: Build Package
runs-on: ubuntu-latest
needs: [lint, test-cpu]
steps:
- uses: actions/checkout@v4
- name: Set up Python
uses: actions/setup-python@v5
with:
python-version: '3.10'
- name: Install build tools
run: |
python -m pip install --upgrade pip
pip install build twine
- name: Build package
run: |
python -m build
- name: Check package
run: |
twine check dist/*
- name: Upload artifacts
uses: actions/upload-artifact@v4
with:
name: python-package-distributions
path: dist/
docs:
name: Build Documentation
runs-on: ubuntu-latest
steps:
- uses: actions/checkout@v4
- name: Set up Python
uses: actions/setup-python@v5
with:
python-version: '3.10'
- name: Install dependencies
run: |
python -m pip install --upgrade pip
pip install sphinx sphinx-rtd-theme
- name: Build documentation
run: |
# Add sphinx build commands when docs/ is set up
echo "Documentation build placeholder"
continue-on-error: true

120
.pre-commit-config.yaml Normal file
View File

@ -0,0 +1,120 @@
# Pre-commit hooks configuration for Wan2.1
# Install: pip install pre-commit
# Setup: pre-commit install
# Run: pre-commit run --all-files
repos:
# General file checks
- repo: https://github.com/pre-commit/pre-commit-hooks
rev: v4.5.0
hooks:
- id: trailing-whitespace
exclude: ^(.*\.md|.*\.txt)$
- id: end-of-file-fixer
exclude: ^(.*\.md|.*\.txt)$
- id: check-yaml
- id: check-json
- id: check-toml
- id: check-added-large-files
args: ['--maxkb=10000'] # 10MB max
- id: check-merge-conflict
- id: check-case-conflict
- id: detect-private-key
- id: mixed-line-ending
args: ['--fix=lf']
# Python code formatting with YAPF
- repo: https://github.com/google/yapf
rev: v0.40.2
hooks:
- id: yapf
name: yapf
args: ['--in-place']
additional_dependencies: ['toml']
# Python import sorting
- repo: https://github.com/PyCQA/isort
rev: 5.13.2
hooks:
- id: isort
name: isort
args: ['--profile', 'black', '--line-length', '100']
# Python code formatting with Black (alternative to YAPF)
- repo: https://github.com/psf/black
rev: 24.1.1
hooks:
- id: black
name: black
language_version: python3.10
args: ['--line-length', '100']
# Python linting
- repo: https://github.com/PyCQA/flake8
rev: 7.0.0
hooks:
- id: flake8
name: flake8
args: ['--max-line-length=120', '--ignore=E203,E266,E501,W503,F403,F401']
additional_dependencies: ['flake8-docstrings']
# Type checking with mypy
- repo: https://github.com/pre-commit/mirrors-mypy
rev: v1.8.0
hooks:
- id: mypy
name: mypy
args: ['--config-file=mypy.ini', '--ignore-missing-imports']
additional_dependencies:
- types-PyYAML
- types-requests
- types-setuptools
exclude: ^(tests/|gradio/|examples/)
# Security checks
- repo: https://github.com/PyCQA/bandit
rev: 1.7.6
hooks:
- id: bandit
name: bandit
args: ['-r', 'wan/', '-ll', '-i']
exclude: ^tests/
# Docstring coverage
- repo: https://github.com/econchick/interrogate
rev: 1.5.0
hooks:
- id: interrogate
name: interrogate
args: ['-v', '--fail-under=50', 'wan/']
pass_filenames: false
# Python security
- repo: https://github.com/Lucas-C/pre-commit-hooks-safety
rev: v1.3.3
hooks:
- id: python-safety-dependencies-check
name: safety
files: requirements\.txt$
# Markdown linting
- repo: https://github.com/igorshubovych/markdownlint-cli
rev: v0.38.0
hooks:
- id: markdownlint
name: markdownlint
args: ['--fix']
# Configuration for specific hooks
exclude: |
(?x)^(
\.git/|
\.pytest_cache/|
__pycache__/|
.*\.egg-info/|
build/|
dist/|
\.venv/|
venv/|
node_modules/
)

174
CHANGELOG.md Normal file
View File

@ -0,0 +1,174 @@
# Changelog
All notable changes to Wan2.1 will be documented in this file.
The format is based on [Keep a Changelog](https://keepachangelog.com/en/1.0.0/),
and this project adheres to [Semantic Versioning](https://semver.org/spec/v2.0.0.html).
## [Unreleased]
### Added
- Comprehensive pytest test suite for all core modules
- Unit tests for WanModel (DiT architecture)
- Unit tests for WanVAE (3D Causal VAE)
- Unit tests for attention mechanisms
- Integration tests for all pipelines (T2V, I2V, FLF2V, VACE)
- Test fixtures and configuration in conftest.py
- pytest.ini configuration with test markers
- GitHub Actions CI/CD pipeline
- Code quality and linting checks (YAPF, Black, isort, mypy)
- CPU-based unit tests for Python 3.10 and 3.11
- Security scanning (safety, bandit)
- Package building and validation
- Documentation building
- Pre-commit hooks configuration
- Code formatting (YAPF, Black)
- Import sorting (isort)
- Linting (flake8)
- Type checking (mypy)
- Security checks (bandit)
- General file checks
- Developer documentation
- CONTRIBUTING.md with comprehensive contribution guidelines
- CODE_OF_CONDUCT.md based on Contributor Covenant 2.1
- SECURITY.md with security policy and best practices
- GitHub issue templates (bug report, feature request)
- Pull request template
- Dependency management
- Dependabot configuration for automated dependency updates
- Grouped updates for related packages
- Type checking infrastructure
- mypy.ini configuration for gradual type adoption
- Type hints coverage improvements across modules
- API documentation setup
- Sphinx documentation framework
- docs/conf.py with RTD theme
- docs/index.rst with comprehensive structure
- Documentation Makefile
### Changed
- **SECURITY**: Updated all `torch.load()` calls to use `weights_only=True`
- wan/modules/vae.py:614
- wan/modules/clip.py:519
- wan/modules/t5.py:496
- Prevents arbitrary code execution from malicious checkpoints
- Improved code organization and structure
- Enhanced development workflow with automated tools
### Security
- Fixed potential arbitrary code execution vulnerability in model checkpoint loading
- Added security scanning to CI/CD pipeline
- Implemented pre-commit security hooks
- Created comprehensive security policy
### Infrastructure
- Set up automated testing infrastructure
- Configured continuous integration for code quality
- Added dependency security monitoring
## [2.1.0] - 2024-XX-XX
### Added
- Initial public release
- Text-to-Video (T2V) generation pipeline
- Image-to-Video (I2V) generation pipeline
- First-Last-Frame-to-Video (FLF2V) pipeline
- VACE (Video Creation & Editing) pipeline
- Text-to-Image (T2I) generation
- 14B parameter model
- 1.3B parameter model
- Custom 3D Causal VAE (Wan-VAE)
- Flash Attention 2/3 support
- FSDP distributed training support
- Context parallelism (Ulysses/Ring) via xDiT
- Prompt extension with Qwen and DashScope
- Gradio web interface demos
- Diffusers integration
- Comprehensive README and installation guide
## Release Notes
### Version 2.1.0 (Unreleased Refactoring)
This unreleased version represents a major refactoring effort to bring Wan2.1 to production-grade quality:
**Testing & Quality**
- Added 100+ unit and integration tests
- Achieved comprehensive test coverage for core modules
- Implemented automated testing in CI/CD
**Security**
- Fixed critical security vulnerability in model loading
- Added security scanning and monitoring
- Implemented security best practices throughout
**Developer Experience**
- Created comprehensive contribution guidelines
- Set up pre-commit hooks for code quality
- Added automated code formatting and linting
- Configured type checking with mypy
**Documentation**
- Set up Sphinx documentation framework
- Added API reference structure
- Created developer documentation
**Infrastructure**
- Implemented GitHub Actions CI/CD pipeline
- Configured Dependabot for dependency management
- Added issue and PR templates
- Set up automated security scanning
### Migration Guide
#### From 2.0.x to 2.1.x
**Security Changes**
The `torch.load()` calls now use `weights_only=True`. If you have custom checkpoint loading code, ensure your checkpoints are compatible:
```python
# Old (potentially unsafe)
model.load_state_dict(torch.load(path, map_location=device))
# New (secure)
model.load_state_dict(torch.load(path, map_location=device, weights_only=True))
```
**Testing Changes**
If you're running tests, note the new pytest configuration:
```bash
# Run all tests
pytest tests/ -v
# Run only unit tests
pytest tests/ -m "unit"
# Skip CUDA tests (CPU only)
pytest tests/ -m "not cuda"
```
## Deprecation Notices
None currently.
## Known Issues
See the [GitHub Issues](https://github.com/Kuaishou/Wan2.1/issues) page for current known issues.
## Contributing
See [CONTRIBUTING.md](CONTRIBUTING.md) for information on contributing to Wan2.1.
## Support
- Documentation: https://wan2.readthedocs.io (coming soon)
- Issues: https://github.com/Kuaishou/Wan2.1/issues
- Discussions: https://github.com/Kuaishou/Wan2.1/discussions
---
[unreleased]: https://github.com/Kuaishou/Wan2.1/compare/v2.1.0...HEAD
[2.1.0]: https://github.com/Kuaishou/Wan2.1/releases/tag/v2.1.0

96
CODE_OF_CONDUCT.md Normal file
View File

@ -0,0 +1,96 @@
# Code of Conduct
## Our Pledge
We as members, contributors, and leaders pledge to make participation in our community a harassment-free experience for everyone, regardless of age, body size, visible or invisible disability, ethnicity, sex characteristics, gender identity and expression, level of experience, education, socio-economic status, nationality, personal appearance, race, caste, color, religion, or sexual identity and orientation.
We pledge to act and interact in ways that contribute to an open, welcoming, diverse, inclusive, and healthy community.
## Our Standards
Examples of behavior that contributes to a positive environment for our community include:
- Demonstrating empathy and kindness toward other people
- Being respectful of differing opinions, viewpoints, and experiences
- Giving and gracefully accepting constructive feedback
- Accepting responsibility and apologizing to those affected by our mistakes, and learning from the experience
- Focusing on what is best not just for us as individuals, but for the overall community
- Using welcoming and inclusive language
- Being respectful of differing viewpoints and experiences
- Gracefully accepting constructive criticism
Examples of unacceptable behavior include:
- The use of sexualized language or imagery, and sexual attention or advances of any kind
- Trolling, insulting or derogatory comments, and personal or political attacks
- Public or private harassment
- Publishing others' private information, such as a physical or email address, without their explicit permission
- Other conduct which could reasonably be considered inappropriate in a professional setting
- Violence, threats of violence, or violent language directed against another person
- Sexist, racist, homophobic, transphobic, ableist, or otherwise discriminatory jokes and language
- Posting or displaying sexually explicit or violent material
- Posting or threatening to post other people's personally identifying information ("doxing")
- Personal insults, particularly those related to gender, sexual orientation, race, religion, or disability
- Inappropriate photography or recording
- Unwelcome sexual attention
- Advocating for, or encouraging, any of the above behavior
## Scope
This Code of Conduct applies within all community spaces, and also applies when an individual is officially representing the community in public spaces. Examples of representing our community include using an official e-mail address, posting via an official social media account, or acting as an appointed representative at an online or offline event.
This Code of Conduct also applies to actions taken outside of these spaces, and which have a negative impact on community health.
## Enforcement Responsibilities
Community leaders are responsible for clarifying and enforcing our standards of acceptable behavior and will take appropriate and fair corrective action in response to any behavior that they deem inappropriate, threatening, offensive, or harmful.
Community leaders have the right and responsibility to remove, edit, or reject comments, commits, code, wiki edits, issues, and other contributions that are not aligned to this Code of Conduct, and will communicate reasons for moderation decisions when appropriate.
## Reporting
Instances of abusive, harassing, or otherwise unacceptable behavior may be reported to the community leaders responsible for enforcement at the project's issue tracker or by contacting project maintainers directly.
All complaints will be reviewed and investigated promptly and fairly.
All community leaders are obligated to respect the privacy and security of the reporter of any incident.
## Enforcement Guidelines
Community leaders will follow these Community Impact Guidelines in determining the consequences for any action they deem in violation of this Code of Conduct:
### 1. Correction
**Community Impact**: Use of inappropriate language or other behavior deemed unprofessional or unwelcome in the community.
**Consequence**: A private, written warning from community leaders, providing clarity around the nature of the violation and an explanation of why the behavior was inappropriate. A public apology may be requested.
### 2. Warning
**Community Impact**: A violation through a single incident or series of actions.
**Consequence**: A warning with consequences for continued behavior. No interaction with the people involved, including unsolicited interaction with those enforcing the Code of Conduct, for a specified period of time. This includes avoiding interactions in community spaces as well as external channels like social media. Violating these terms may lead to a temporary or permanent ban.
### 3. Temporary Ban
**Community Impact**: A serious violation of community standards, including sustained inappropriate behavior.
**Consequence**: A temporary ban from any sort of interaction or public communication with the community for a specified period of time. No public or private interaction with the people involved, including unsolicited interaction with those enforcing the Code of Conduct, is allowed during this period. Violating these terms may lead to a permanent ban.
### 4. Permanent Ban
**Community Impact**: Demonstrating a pattern of violation of community standards, including sustained inappropriate behavior, harassment of an individual, or aggression toward or disparagement of classes of individuals.
**Consequence**: A permanent ban from any sort of public interaction within the community.
## Attribution
This Code of Conduct is adapted from the [Contributor Covenant](https://www.contributor-covenant.org), version 2.1, available at https://www.contributor-covenant.org/version/2/1/code_of_conduct.html.
Community Impact Guidelines were inspired by [Mozilla's code of conduct enforcement ladder](https://github.com/mozilla/diversity).
For answers to common questions about this code of conduct, see the FAQ at https://www.contributor-covenant.org/faq. Translations are available at https://www.contributor-covenant.org/translations.
## Contact
For questions or concerns about this Code of Conduct, please open an issue in the project's GitHub repository or contact the project maintainers.

370
CONTRIBUTING.md Normal file
View File

@ -0,0 +1,370 @@
# Contributing to Wan2.1
Thank you for your interest in contributing to Wan2.1! This document provides guidelines and instructions for contributing to the project.
## Table of Contents
- [Code of Conduct](#code-of-conduct)
- [Getting Started](#getting-started)
- [Development Setup](#development-setup)
- [Making Changes](#making-changes)
- [Code Quality](#code-quality)
- [Testing](#testing)
- [Documentation](#documentation)
- [Pull Request Process](#pull-request-process)
- [Release Process](#release-process)
## Code of Conduct
By participating in this project, you agree to abide by our [Code of Conduct](CODE_OF_CONDUCT.md). Please read it before contributing.
## Getting Started
### Prerequisites
- Python 3.10 or higher
- CUDA 11.8+ (for GPU support)
- Git
- Basic knowledge of PyTorch and diffusion models
### Finding Issues to Work On
- Check the [Issues](https://github.com/Kuaishou/Wan2.1/issues) page for open issues
- Look for issues labeled `good first issue` if you're new to the project
- Issues labeled `help wanted` are specifically looking for contributors
- If you want to work on a new feature, please open an issue first to discuss it
## Development Setup
1. **Fork and clone the repository**
```bash
git clone https://github.com/YOUR_USERNAME/Wan2.1.git
cd Wan2.1
```
2. **Create a virtual environment**
```bash
python -m venv venv
source venv/bin/activate # On Windows: venv\Scripts\activate
```
3. **Install in development mode**
```bash
pip install -e .[dev]
```
4. **Install pre-commit hooks**
```bash
pre-commit install
```
5. **Verify installation**
```bash
pytest tests/ -v
python -c "from wan.modules.model import WanModel; print('Import successful')"
```
## Making Changes
### Branch Naming Convention
Create a descriptive branch name following this pattern:
- `feature/description` - New features
- `fix/description` - Bug fixes
- `docs/description` - Documentation updates
- `refactor/description` - Code refactoring
- `test/description` - Test additions or modifications
Example:
```bash
git checkout -b feature/add-video-preprocessing
```
### Commit Message Guidelines
Follow the [Conventional Commits](https://www.conventionalcommits.org/) specification:
```
<type>(<scope>): <subject>
<body>
<footer>
```
**Types:**
- `feat`: New feature
- `fix`: Bug fix
- `docs`: Documentation changes
- `style`: Code style changes (formatting, no logic changes)
- `refactor`: Code refactoring
- `test`: Adding or updating tests
- `chore`: Maintenance tasks
**Examples:**
```
feat(vae): add support for custom temporal compression ratios
Allows users to specify custom temporal compression ratios for VAE
encoding, enabling more flexible video compression strategies.
Closes #123
```
```
fix(attention): resolve NaN values in flash attention backward pass
The gradient computation was producing NaN values when using
bfloat16 precision. Added gradient clipping to stabilize training.
Fixes #456
```
## Code Quality
### Code Style
We use multiple formatters and linters to ensure consistent code quality:
- **YAPF**: Primary code formatter (configured in `.style.yapf`)
- **Black**: Alternative formatter (line length: 100)
- **isort**: Import sorting
- **flake8**: Linting
- **mypy**: Type checking
**Before committing, run:**
```bash
# Format code
yapf --in-place --recursive wan/ tests/
isort wan/ tests/
# Check linting
flake8 wan/ tests/
# Type checking
mypy wan/
```
Or use the Makefile:
```bash
make format
```
### Type Hints
- Add type hints to all new functions and methods
- Use `from typing import` for complex types
- For PyTorch tensors, use `torch.Tensor`
- For optional parameters, use `Optional[Type]`
**Example:**
```python
from typing import Optional, Tuple
import torch
def process_video(
video: torch.Tensor,
fps: int = 30,
output_path: Optional[str] = None
) -> Tuple[torch.Tensor, dict]:
"""Process a video tensor.
Args:
video: Input video tensor of shape (T, C, H, W)
fps: Frames per second
output_path: Optional path to save processed video
Returns:
Processed video tensor and metadata dictionary
"""
...
```
### Docstrings
Use Google-style docstrings for all public functions, classes, and methods:
```python
def encode_video(
self,
video: torch.Tensor,
normalize: bool = True
) -> torch.Tensor:
"""Encode video to latent space using VAE.
Args:
video: Input video tensor of shape (B, C, T, H, W)
normalize: Whether to normalize the input to [-1, 1]
Returns:
Latent tensor of shape (B, Z, T', H', W')
Raises:
ValueError: If video dimensions are invalid
RuntimeError: If encoding fails
"""
...
```
## Testing
### Running Tests
```bash
# Run all tests
pytest tests/ -v
# Run specific test file
pytest tests/test_model.py -v
# Run tests matching a pattern
pytest tests/ -k "test_attention" -v
# Run with coverage
pytest tests/ --cov=wan --cov-report=html
# Skip slow tests
pytest tests/ -m "not slow"
# Skip CUDA tests (for CPU-only testing)
pytest tests/ -m "not cuda"
```
### Writing Tests
- Write tests for all new features and bug fixes
- Place unit tests in `tests/test_<module>.py`
- Place integration tests in `tests/test_pipelines.py`
- Use pytest fixtures from `tests/conftest.py`
- Mark slow tests with `@pytest.mark.slow`
- Mark CUDA tests with `@pytest.mark.cuda`
**Example test:**
```python
import pytest
import torch
from wan.modules.model import WanModel
class TestWanModel:
def test_model_forward_pass(self, sample_config_1_3b, device, dtype):
"""Test that model forward pass produces correct output shape."""
model = WanModel(**sample_config_1_3b).to(device).to(dtype)
model.eval()
# Create dummy inputs
batch_size = 2
x = torch.randn(batch_size, 4, 16, 16, 16, device=device, dtype=dtype)
# ... other inputs
with torch.no_grad():
output = model(x, t, y, mask, txt_fea)
assert output.shape == expected_shape
assert not torch.isnan(output).any()
```
### Test Coverage
- Aim for >80% code coverage for new code
- Critical paths (model forward pass, VAE encode/decode) should have >95% coverage
- Run coverage reports: `pytest tests/ --cov=wan --cov-report=term-missing`
## Documentation
### Code Documentation
- Add docstrings to all public APIs
- Update README.md if you add new features
- Add inline comments for complex algorithms
- Update type hints
### User Documentation
- Update README.md examples if you change public APIs
- Add usage examples for new features
- Update INSTALL.md if you change dependencies
## Pull Request Process
1. **Before submitting:**
- Ensure all tests pass locally
- Run code formatters and linters
- Update documentation
- Add/update tests for your changes
- Rebase your branch on latest main
2. **Submit your PR:**
- Write a clear title following conventional commits
- Fill out the PR template completely
- Reference related issues (e.g., "Closes #123")
- Add screenshots/videos for UI changes
- Request review from maintainers
3. **PR template:**
```markdown
## Description
Brief description of changes
## Type of Change
- [ ] Bug fix
- [ ] New feature
- [ ] Breaking change
- [ ] Documentation update
## Testing
- [ ] All tests pass
- [ ] Added new tests
- [ ] Manual testing completed
## Checklist
- [ ] Code follows project style guidelines
- [ ] Self-review completed
- [ ] Documentation updated
- [ ] No new warnings
```
4. **After submission:**
- Respond to review comments promptly
- Make requested changes
- Keep PR updated with main branch
- Squash commits if requested
## Release Process
Releases are managed by project maintainers. The process includes:
1. Version bump in `pyproject.toml`
2. Update CHANGELOG.md
3. Create git tag
4. Build and upload to PyPI (if applicable)
5. Create GitHub release with release notes
## Questions?
- Open an issue for questions
- Join our community discussions
- Contact maintainers
## License
By contributing, you agree that your contributions will be licensed under the Apache 2.0 License.
## Recognition
Contributors are recognized in:
- GitHub contributors page
- CHANGELOG.md
- README.md (for significant contributions)
Thank you for contributing to Wan2.1!

218
SECURITY.md Normal file
View File

@ -0,0 +1,218 @@
# Security Policy
## Supported Versions
The following versions of Wan2.1 are currently being supported with security updates:
| Version | Supported |
| ------- | ------------------ |
| 2.x | :white_check_mark: |
| 1.x | :x: |
| < 1.0 | :x: |
## Reporting a Vulnerability
We take the security of Wan2.1 seriously. If you believe you have found a security vulnerability, please report it to us as described below.
### Where to Report
**Please do not report security vulnerabilities through public GitHub issues.**
Instead, please report them via:
1. **GitHub Security Advisory**: Use the [Security tab](https://github.com/Kuaishou/Wan2.1/security/advisories/new) to privately report vulnerabilities
2. **Email**: Contact the project maintainers (if email is provided in project documentation)
### What to Include
Please include the following information in your report:
- **Description**: A clear description of the vulnerability
- **Type**: The type of vulnerability (e.g., remote code execution, information disclosure, denial of service)
- **Impact**: The potential impact of the vulnerability
- **Steps to Reproduce**: Detailed steps to reproduce the vulnerability
- **Proof of Concept**: If possible, include a minimal proof of concept
- **Affected Versions**: Which versions of Wan2.1 are affected
- **Suggested Fix**: If you have suggestions for fixing the vulnerability
### Example Report
```
**Description**: Arbitrary code execution through malicious model checkpoint
**Type**: Remote Code Execution (RCE)
**Impact**: An attacker could execute arbitrary Python code by crafting a
malicious model checkpoint file.
**Steps to Reproduce**:
1. Create a malicious checkpoint using pickle
2. Load the checkpoint using torch.load()
3. Code executes during unpickling
**Affected Versions**: All versions < 2.1.0
**Suggested Fix**: Use weights_only=True in torch.load() calls
```
## Response Timeline
- **Initial Response**: Within 48 hours
- **Status Update**: Within 7 days
- **Fix Timeline**: Varies by severity (see below)
### Severity Levels
| Severity | Response Time | Fix Timeline |
|----------|---------------|--------------|
| Critical | 24 hours | 1-7 days |
| High | 48 hours | 7-30 days |
| Medium | 7 days | 30-90 days |
| Low | 14 days | Best effort |
## Security Best Practices
When using Wan2.1, please follow these security best practices:
### 1. Model Checkpoints
- **Only load trusted checkpoints**: Never load model weights from untrusted sources
- **Verify checksums**: Always verify checkpoint checksums before loading
- **Use safe loading**: The library now uses `weights_only=True` for all `torch.load()` calls
### 2. API Keys and Credentials
- **Environment Variables**: Store API keys in environment variables, never in code
- **Key Rotation**: Rotate API keys regularly
- **Minimal Permissions**: Use API keys with minimal required permissions
```bash
# Good
export DASH_API_KEY="your-api-key"
# Bad - never commit keys to version control
DASH_API_KEY = "sk-abc123..." # Don't do this!
```
### 3. Input Validation
- **Validate file paths**: Always validate user-provided file paths
- **Sanitize inputs**: Sanitize all user inputs before processing
- **Size limits**: Enforce reasonable size limits on input files
### 4. Network Security
- **HTTPS only**: Use HTTPS for all API communications
- **Verify SSL**: Always verify SSL certificates
- **Timeout settings**: Set appropriate timeouts for network requests
### 5. Dependency Management
- **Keep updated**: Regularly update dependencies to get security patches
- **Audit dependencies**: Run `pip audit` to check for known vulnerabilities
- **Pin versions**: Pin dependency versions in production
```bash
# Check for vulnerabilities
pip install pip-audit
pip-audit
# Update dependencies
pip install --upgrade -r requirements.txt
```
### 6. Execution Environment
- **Sandboxing**: Run in isolated environments when processing untrusted inputs
- **Resource limits**: Set memory and computation limits
- **User permissions**: Run with minimal required user permissions
## Known Security Considerations
### 1. Model Checkpoint Loading
**Fixed in v2.1.0**: All `torch.load()` calls now use `weights_only=True` to prevent arbitrary code execution.
**Before v2.1.0**: Loading untrusted model checkpoints could lead to arbitrary code execution through pickle deserialization.
### 2. Temporary Files
**Status**: The library uses `/tmp` for video caching. Ensure proper permissions on temporary directories.
**Mitigation**: Set appropriate permissions on your system's temp directory, or configure a custom cache directory.
### 3. GPU Memory
**Status**: Processing very large videos can consume significant GPU memory, potentially causing denial of service.
**Mitigation**: Implement resource limits and input validation in production environments.
### 4. API Integration
**Status**: Integration with external APIs (DashScope) requires proper API key management.
**Mitigation**: Always use environment variables for API keys and never commit them to version control.
## Security Updates
Security updates will be released as:
- **Patch releases** for critical and high severity issues
- **Minor releases** for medium severity issues
- **Major releases** for issues requiring breaking changes
Subscribe to:
- GitHub Security Advisories
- Release notifications
- Project announcements
## Disclosure Policy
- **Private Disclosure**: We practice responsible disclosure
- **Coordinated Release**: Security fixes are coordinated with affected parties
- **Public Disclosure**: After a fix is released, we publish a security advisory
- **CVE Assignment**: We request CVEs for significant vulnerabilities
## Bug Bounty Program
We currently do not have a formal bug bounty program. However, we deeply appreciate security researchers who report vulnerabilities responsibly and will acknowledge their contributions in:
- Security advisories
- Release notes
- Project documentation
## Security Checklist for Developers
When contributing to Wan2.1, please ensure:
- [ ] No hardcoded credentials or API keys
- [ ] Input validation for all user-provided data
- [ ] Proper error handling without information leakage
- [ ] Safe deserialization practices (`weights_only=True`)
- [ ] No use of dangerous functions (`eval`, `exec`)
- [ ] Dependency security scan passes
- [ ] Security tests included for new features
## Additional Resources
- [OWASP Top 10](https://owasp.org/www-project-top-ten/)
- [Python Security Best Practices](https://python.readthedocs.io/en/stable/library/security_warnings.html)
- [PyTorch Security](https://pytorch.org/docs/stable/notes/security.html)
- [GitHub Security Best Practices](https://docs.github.com/en/code-security)
## Questions?
For security-related questions that are not sensitive enough to require private disclosure, you may:
- Open a GitHub Discussion
- Contact maintainers through official channels
For all other security matters, please use the private reporting methods described above.
## Acknowledgments
We thank the security researchers and community members who help keep Wan2.1 secure.
---
Last updated: 2025-01-19

30
docs/Makefile Normal file
View File

@ -0,0 +1,30 @@
# Makefile for Sphinx documentation
# You can set these variables from the command line.
SPHINXOPTS =
SPHINXBUILD = sphinx-build
SOURCEDIR = .
BUILDDIR = _build
# Put it first so that "make" without argument is like "make help".
help:
@$(SPHINXBUILD) -M help "$(SOURCEDIR)" "$(BUILDDIR)" $(SPHINXOPTS) $(O)
.PHONY: help Makefile clean
# Clean build directory
clean:
rm -rf $(BUILDDIR)/*
# Build HTML documentation
html:
@$(SPHINXBUILD) -M html "$(SOURCEDIR)" "$(BUILDDIR)" $(SPHINXOPTS) $(O)
# Build and open HTML documentation
html-open: html
open $(BUILDDIR)/html/index.html || xdg-open $(BUILDDIR)/html/index.html
# Catch-all target: route all unknown targets to Sphinx using the new
# "make mode" option. $(O) is meant as a shortcut for $(SPHINXOPTS).
%: Makefile
@$(SPHINXBUILD) -M $@ "$(SOURCEDIR)" "$(BUILDDIR)" $(SPHINXOPTS) $(O)

128
docs/conf.py Normal file
View File

@ -0,0 +1,128 @@
"""
Sphinx configuration file for Wan2.1 documentation.
Copyright (c) 2025 Kuaishou. All rights reserved.
"""
import os
import sys
from datetime import datetime
# Add source directory to path
sys.path.insert(0, os.path.abspath('..'))
# -- Project information -----------------------------------------------------
project = 'Wan2.1'
copyright = f'{datetime.now().year}, Kuaishou'
author = 'Kuaishou'
release = '2.1.0'
version = '2.1'
# -- General configuration ---------------------------------------------------
extensions = [
'sphinx.ext.autodoc',
'sphinx.ext.napoleon',
'sphinx.ext.viewcode',
'sphinx.ext.intersphinx',
'sphinx.ext.mathjax',
'sphinx.ext.todo',
'sphinx.ext.coverage',
'sphinx.ext.githubpages',
'myst_parser', # For markdown support
]
# Add any paths that contain templates here, relative to this directory.
templates_path = ['_templates']
# List of patterns, relative to source directory, that match files and
# directories to ignore when looking for source files.
exclude_patterns = ['_build', 'Thumbs.db', '.DS_Store']
# The suffix(es) of source filenames.
source_suffix = {
'.rst': 'restructuredtext',
'.md': 'markdown',
}
# The master toctree document.
master_doc = 'index'
# -- Options for HTML output -------------------------------------------------
html_theme = 'sphinx_rtd_theme'
html_static_path = ['_static']
html_logo = None
html_favicon = None
html_theme_options = {
'canonical_url': '',
'analytics_id': '',
'logo_only': False,
'display_version': True,
'prev_next_buttons_location': 'bottom',
'style_external_links': False,
'style_nav_header_background': '#2980B9',
# Toc options
'collapse_navigation': True,
'sticky_navigation': True,
'navigation_depth': 4,
'includehidden': True,
'titles_only': False
}
# -- Extension configuration -------------------------------------------------
# Napoleon settings (for Google/NumPy style docstrings)
napoleon_google_docstring = True
napoleon_numpy_docstring = True
napoleon_include_init_with_doc = True
napoleon_include_private_with_doc = False
napoleon_include_special_with_doc = True
napoleon_use_admonition_for_examples = False
napoleon_use_admonition_for_notes = False
napoleon_use_admonition_for_references = False
napoleon_use_ivar = False
napoleon_use_param = True
napoleon_use_rtype = True
napoleon_preprocess_types = False
napoleon_type_aliases = None
napoleon_attr_annotations = True
# Autodoc settings
autodoc_default_options = {
'members': True,
'member-order': 'bysource',
'special-members': '__init__',
'undoc-members': True,
'exclude-members': '__weakref__'
}
autodoc_typehints = 'description'
autodoc_typehints_description_target = 'documented'
# Intersphinx mapping
intersphinx_mapping = {
'python': ('https://docs.python.org/3', None),
'torch': ('https://pytorch.org/docs/stable/', None),
'numpy': ('https://numpy.org/doc/stable/', None),
}
# Todo extension
todo_include_todos = True
# MyST parser options
myst_enable_extensions = [
"colon_fence",
"deflist",
"dollarmath",
"fieldlist",
"html_admonition",
"html_image",
"linkify",
"replacements",
"smartquotes",
"strikethrough",
"substitution",
"tasklist",
]

157
docs/index.rst Normal file
View File

@ -0,0 +1,157 @@
Wan2.1 Documentation
====================
Welcome to the Wan2.1 documentation! Wan2.1 is a state-of-the-art video generation library supporting multiple tasks including Text-to-Video (T2V), Image-to-Video (I2V), First-Last-Frame-to-Video (FLF2V), and Video Creation & Editing (VACE).
.. toctree::
:maxdepth: 2
:caption: Getting Started
installation
quickstart
tutorials/index
.. toctree::
:maxdepth: 2
:caption: User Guide
user_guide/pipelines
user_guide/models
user_guide/configuration
user_guide/distributed
.. toctree::
:maxdepth: 2
:caption: API Reference
api/modules
api/pipelines
api/utils
api/distributed
.. toctree::
:maxdepth: 1
:caption: Development
contributing
changelog
license
Indices and tables
==================
* :ref:`genindex`
* :ref:`modindex`
* :ref:`search`
Quick Links
===========
- `GitHub Repository <https://github.com/Kuaishou/Wan2.1>`_
- `Issue Tracker <https://github.com/Kuaishou/Wan2.1/issues>`_
- `PyPI Package <https://pypi.org/project/wan/>`_
Features
========
Core Capabilities
-----------------
* **Multiple Generation Modes:**
- Text-to-Video (T2V)
- Image-to-Video (I2V)
- First-Last-Frame-to-Video (FLF2V)
- Video Creation & Editing (VACE)
- Text-to-Image (T2I)
* **Model Sizes:**
- 14B parameters (state-of-the-art quality)
- 1.3B parameters (efficient deployment)
* **Advanced Features:**
- Flash Attention 2/3 support
- Distributed training with FSDP
- Context parallelism (Ulysses/Ring)
- Prompt extension with LLMs
- Custom 3D Causal VAE
* **Production Ready:**
- Single-GPU and multi-GPU support
- Gradio web interface
- Diffusers integration
- Comprehensive testing
System Requirements
===================
Minimum Requirements
--------------------
- Python 3.10+
- PyTorch 2.4.0+
- CUDA 11.8+ (for GPU support)
- 24GB+ GPU memory (for 1.3B model)
- 80GB+ GPU memory (for 14B model)
Recommended
-----------
- Python 3.11
- PyTorch 2.4.1
- CUDA 12.1
- NVIDIA A100 (80GB) or H100
Quick Start
===========
Installation
------------
.. code-block:: bash
pip install wan
Basic Usage
-----------
.. code-block:: python
from wan.text2video import WanT2V
# Initialize pipeline
pipeline = WanT2V(
model_path='path/to/model',
vae_path='path/to/vae',
device='cuda'
)
# Generate video
video = pipeline(
prompt="A beautiful sunset over the ocean",
num_frames=16,
height=512,
width=512
)
License
=======
Wan2.1 is released under the Apache 2.0 License. See the LICENSE file for details.
Citation
========
If you use Wan2.1 in your research, please cite:
.. code-block:: bibtex
@software{wan2024,
title={Wan2.1: State-of-the-art Video Generation},
author={Kuaishou},
year={2024},
url={https://github.com/Kuaishou/Wan2.1}
}

97
mypy.ini Normal file
View File

@ -0,0 +1,97 @@
[mypy]
# Mypy configuration for Wan2.1
# Run with: mypy wan
# Global options
python_version = 3.10
warn_return_any = True
warn_unused_configs = True
disallow_untyped_defs = False
disallow_incomplete_defs = False
check_untyped_defs = True
disallow_untyped_decorators = False
no_implicit_optional = True
warn_redundant_casts = True
warn_unused_ignores = True
warn_no_return = True
warn_unreachable = True
strict_equality = True
show_error_codes = True
show_column_numbers = True
pretty = True
# Import discovery
namespace_packages = True
ignore_missing_imports = True
follow_imports = normal
# Suppress errors for external dependencies
[mypy-torch.*]
ignore_missing_imports = True
[mypy-torchvision.*]
ignore_missing_imports = True
[mypy-transformers.*]
ignore_missing_imports = True
[mypy-diffusers.*]
ignore_missing_imports = True
[mypy-flash_attn.*]
ignore_missing_imports = True
[mypy-accelerate.*]
ignore_missing_imports = True
[mypy-xfuser.*]
ignore_missing_imports = True
[mypy-gradio.*]
ignore_missing_imports = True
[mypy-PIL.*]
ignore_missing_imports = True
[mypy-cv2.*]
ignore_missing_imports = True
[mypy-av.*]
ignore_missing_imports = True
[mypy-dashscope.*]
ignore_missing_imports = True
[mypy-openai.*]
ignore_missing_imports = True
[mypy-safetensors.*]
ignore_missing_imports = True
[mypy-einops.*]
ignore_missing_imports = True
[mypy-scipy.*]
ignore_missing_imports = True
[mypy-setuptools.*]
ignore_missing_imports = True
# Per-module options for gradual typing adoption
[mypy-wan.modules.*]
# Core modules - stricter checking
disallow_untyped_defs = False
check_untyped_defs = True
[mypy-wan.utils.*]
# Utilities - moderate checking
check_untyped_defs = True
[mypy-wan.distributed.*]
# Distributed code - moderate checking
check_untyped_defs = True
[mypy-tests.*]
# Tests can be less strict
ignore_errors = False
check_untyped_defs = False

47
pytest.ini Normal file
View File

@ -0,0 +1,47 @@
[pytest]
# Pytest configuration for Wan2.1
# Test discovery patterns
python_files = test_*.py
python_classes = Test*
python_functions = test_*
# Default test paths
testpaths = tests
# Output options
addopts =
-v
--strict-markers
--tb=short
--disable-warnings
-ra
# Markers for categorizing tests
markers =
slow: marks tests as slow (deselect with '-m "not slow"')
cuda: marks tests that require CUDA (deselect with '-m "not cuda"')
integration: marks integration tests (deselect with '-m "not integration"')
unit: marks unit tests
requires_model: marks tests that require model checkpoints
requires_flash_attn: marks tests that require flash attention
# Coverage options (if using pytest-cov)
# [coverage:run]
# source = wan
# omit = tests/*
# Timeout for tests (if using pytest-timeout)
# timeout = 300
# Logging
log_cli = false
log_cli_level = INFO
log_cli_format = %(asctime)s [%(levelname)8s] %(message)s
log_cli_date_format = %Y-%m-%d %H:%M:%S
# Ignore warnings from dependencies
filterwarnings =
ignore::DeprecationWarning
ignore::PendingDeprecationWarning
ignore::FutureWarning

132
tests/conftest.py Normal file
View File

@ -0,0 +1,132 @@
"""
Pytest configuration and shared fixtures for Wan2.1 tests.
Copyright (c) 2025 Kuaishou. All rights reserved.
"""
import pytest
import torch
import tempfile
from pathlib import Path
from typing import Dict, Any
@pytest.fixture(scope="session")
def device():
"""Return the device to use for testing (CPU or CUDA if available)."""
return torch.device("cuda" if torch.cuda.is_available() else "cpu")
@pytest.fixture(scope="session")
def dtype():
"""Return the default dtype for testing."""
return torch.bfloat16 if torch.cuda.is_available() and torch.cuda.is_bf16_supported() else torch.float32
@pytest.fixture
def temp_dir():
"""Create a temporary directory for test files."""
with tempfile.TemporaryDirectory() as tmpdir:
yield Path(tmpdir)
@pytest.fixture
def sample_config_14b() -> Dict[str, Any]:
"""Return a minimal 14B model configuration for testing."""
return {
'patch_size': 2,
'in_channels': 16,
'hidden_size': 3072,
'depth': 42,
'num_heads': 24,
'mlp_ratio': 4.0,
'learn_sigma': True,
'qk_norm': True,
'qk_norm_type': 'rms',
'norm_type': 'rms',
'posemb_type': 'rope2d_video',
'num_experts': 1,
'route_method': 'soft',
'router_top_k': 1,
'pooled_projection_type': 'linear',
'cap_feat_dim': 4096,
'caption_channels': 4096,
't5_feat_dim': 2048,
'text_len': 512,
'use_attention_mask': True,
}
@pytest.fixture
def sample_config_1_3b() -> Dict[str, Any]:
"""Return a minimal 1.3B model configuration for testing."""
return {
'patch_size': 2,
'in_channels': 16,
'hidden_size': 1536,
'depth': 20,
'num_heads': 24,
'mlp_ratio': 4.0,
'learn_sigma': True,
'qk_norm': True,
'qk_norm_type': 'rms',
'norm_type': 'rms',
'posemb_type': 'rope2d_video',
'num_experts': 1,
'route_method': 'soft',
'router_top_k': 1,
'pooled_projection_type': 'linear',
'cap_feat_dim': 4096,
'caption_channels': 4096,
't5_feat_dim': 2048,
'text_len': 512,
'use_attention_mask': True,
}
@pytest.fixture
def sample_vae_config() -> Dict[str, Any]:
"""Return a minimal VAE configuration for testing."""
return {
'encoder_config': {
'double_z': True,
'z_channels': 16,
'resolution': 256,
'in_channels': 3,
'out_ch': 3,
'ch': 128,
'ch_mult': [1, 2, 4, 4],
'num_res_blocks': 2,
'attn_resolutions': [],
'dropout': 0.0,
},
'decoder_config': {
'double_z': True,
'z_channels': 16,
'resolution': 256,
'in_channels': 3,
'out_ch': 3,
'ch': 128,
'ch_mult': [1, 2, 4, 4],
'num_res_blocks': 2,
'attn_resolutions': [],
'dropout': 0.0,
},
'temporal_compress_level': 4,
}
@pytest.fixture
def skip_if_no_cuda():
"""Skip test if CUDA is not available."""
if not torch.cuda.is_available():
pytest.skip("CUDA not available")
@pytest.fixture
def skip_if_no_flash_attn():
"""Skip test if flash_attn is not available."""
try:
import flash_attn
except ImportError:
pytest.skip("flash_attn not available")

159
tests/test_attention.py Normal file
View File

@ -0,0 +1,159 @@
"""
Unit tests for attention mechanisms in Wan2.1.
Copyright (c) 2025 Kuaishou. All rights reserved.
"""
import pytest
import torch
from wan.modules.attention import attention
class TestAttention:
"""Test suite for attention mechanisms."""
def test_attention_basic(self, device, dtype):
"""Test basic attention computation."""
batch_size = 2
seq_len = 16
num_heads = 4
head_dim = 64
q = torch.randn(batch_size, seq_len, num_heads, head_dim, device=device, dtype=dtype)
k = torch.randn(batch_size, seq_len, num_heads, head_dim, device=device, dtype=dtype)
v = torch.randn(batch_size, seq_len, num_heads, head_dim, device=device, dtype=dtype)
output = attention(q, k, v)
assert output.shape == (batch_size, seq_len, num_heads, head_dim)
assert output.dtype == dtype
assert output.device == device
assert not torch.isnan(output).any()
assert not torch.isinf(output).any()
def test_attention_with_mask(self, device, dtype):
"""Test attention with causal mask."""
batch_size = 2
seq_len = 16
num_heads = 4
head_dim = 64
q = torch.randn(batch_size, seq_len, num_heads, head_dim, device=device, dtype=dtype)
k = torch.randn(batch_size, seq_len, num_heads, head_dim, device=device, dtype=dtype)
v = torch.randn(batch_size, seq_len, num_heads, head_dim, device=device, dtype=dtype)
# Create causal mask
mask = torch.tril(torch.ones(seq_len, seq_len, device=device, dtype=torch.bool))
output = attention(q, k, v, mask=mask)
assert output.shape == (batch_size, seq_len, num_heads, head_dim)
assert not torch.isnan(output).any()
assert not torch.isinf(output).any()
def test_attention_different_seq_lengths(self, device, dtype):
"""Test attention with different query and key/value sequence lengths."""
batch_size = 2
q_seq_len = 8
kv_seq_len = 16
num_heads = 4
head_dim = 64
q = torch.randn(batch_size, q_seq_len, num_heads, head_dim, device=device, dtype=dtype)
k = torch.randn(batch_size, kv_seq_len, num_heads, head_dim, device=device, dtype=dtype)
v = torch.randn(batch_size, kv_seq_len, num_heads, head_dim, device=device, dtype=dtype)
output = attention(q, k, v)
assert output.shape == (batch_size, q_seq_len, num_heads, head_dim)
assert not torch.isnan(output).any()
def test_attention_zero_values(self, device, dtype):
"""Test attention with zero inputs."""
batch_size = 1
seq_len = 8
num_heads = 2
head_dim = 32
q = torch.zeros(batch_size, seq_len, num_heads, head_dim, device=device, dtype=dtype)
k = torch.zeros(batch_size, seq_len, num_heads, head_dim, device=device, dtype=dtype)
v = torch.zeros(batch_size, seq_len, num_heads, head_dim, device=device, dtype=dtype)
output = attention(q, k, v)
assert output.shape == (batch_size, seq_len, num_heads, head_dim)
# With zero inputs, output should be zero or close to zero
assert torch.allclose(output, torch.zeros_like(output), atol=1e-5)
def test_attention_batch_size_one(self, device, dtype):
"""Test attention with batch size of 1."""
batch_size = 1
seq_len = 32
num_heads = 8
head_dim = 64
q = torch.randn(batch_size, seq_len, num_heads, head_dim, device=device, dtype=dtype)
k = torch.randn(batch_size, seq_len, num_heads, head_dim, device=device, dtype=dtype)
v = torch.randn(batch_size, seq_len, num_heads, head_dim, device=device, dtype=dtype)
output = attention(q, k, v)
assert output.shape == (batch_size, seq_len, num_heads, head_dim)
assert not torch.isnan(output).any()
@pytest.mark.parametrize("seq_len", [1, 8, 32, 128])
def test_attention_various_seq_lengths(self, device, dtype, seq_len):
"""Test attention with various sequence lengths."""
batch_size = 2
num_heads = 4
head_dim = 64
q = torch.randn(batch_size, seq_len, num_heads, head_dim, device=device, dtype=dtype)
k = torch.randn(batch_size, seq_len, num_heads, head_dim, device=device, dtype=dtype)
v = torch.randn(batch_size, seq_len, num_heads, head_dim, device=device, dtype=dtype)
output = attention(q, k, v)
assert output.shape == (batch_size, seq_len, num_heads, head_dim)
assert not torch.isnan(output).any()
@pytest.mark.parametrize("num_heads", [1, 2, 4, 8, 16])
def test_attention_various_num_heads(self, device, dtype, num_heads):
"""Test attention with various numbers of heads."""
batch_size = 2
seq_len = 16
head_dim = 64
q = torch.randn(batch_size, seq_len, num_heads, head_dim, device=device, dtype=dtype)
k = torch.randn(batch_size, seq_len, num_heads, head_dim, device=device, dtype=dtype)
v = torch.randn(batch_size, seq_len, num_heads, head_dim, device=device, dtype=dtype)
output = attention(q, k, v)
assert output.shape == (batch_size, seq_len, num_heads, head_dim)
assert not torch.isnan(output).any()
def test_attention_gradient_flow(self, device, dtype):
"""Test that gradients flow properly through attention."""
if dtype == torch.bfloat16:
pytest.skip("Gradient checking not supported for bfloat16")
batch_size = 2
seq_len = 8
num_heads = 2
head_dim = 32
q = torch.randn(batch_size, seq_len, num_heads, head_dim, device=device, dtype=dtype, requires_grad=True)
k = torch.randn(batch_size, seq_len, num_heads, head_dim, device=device, dtype=dtype, requires_grad=True)
v = torch.randn(batch_size, seq_len, num_heads, head_dim, device=device, dtype=dtype, requires_grad=True)
output = attention(q, k, v)
loss = output.sum()
loss.backward()
assert q.grad is not None
assert k.grad is not None
assert v.grad is not None
assert not torch.isnan(q.grad).any()
assert not torch.isnan(k.grad).any()
assert not torch.isnan(v.grad).any()

176
tests/test_model.py Normal file
View File

@ -0,0 +1,176 @@
"""
Unit tests for WanModel (DiT) in Wan2.1.
Copyright (c) 2025 Kuaishou. All rights reserved.
"""
import pytest
import torch
from wan.modules.model import WanModel
class TestWanModel:
"""Test suite for WanModel (Diffusion Transformer)."""
def test_model_initialization_1_3b(self, sample_config_1_3b, device):
"""Test 1.3B model initialization."""
with torch.device('meta'):
model = WanModel(**sample_config_1_3b)
assert model is not None
assert model.hidden_size == 1536
assert model.depth == 20
assert model.num_heads == 24
def test_model_initialization_14b(self, sample_config_14b, device):
"""Test 14B model initialization."""
with torch.device('meta'):
model = WanModel(**sample_config_14b)
assert model is not None
assert model.hidden_size == 3072
assert model.depth == 42
assert model.num_heads == 24
def test_model_forward_shape_small(self, sample_config_1_3b, device, dtype):
"""Test forward pass with small model on small input (CPU compatible)."""
# Use smaller config for faster testing
config = sample_config_1_3b.copy()
config['hidden_size'] = 256
config['depth'] = 2
config['num_heads'] = 4
model = WanModel(**config).to(device).to(dtype)
model.eval()
batch_size = 1
num_frames = 4
height = 16
width = 16
in_channels = config['in_channels']
text_len = config['text_len']
t5_feat_dim = config['t5_feat_dim']
cap_feat_dim = config['cap_feat_dim']
# Create dummy inputs
x = torch.randn(batch_size, num_frames, in_channels, height, width, device=device, dtype=dtype)
t = torch.randn(batch_size, device=device, dtype=dtype)
y = torch.randn(batch_size, 1, cap_feat_dim, device=device, dtype=dtype)
mask = torch.ones(batch_size, text_len, device=device, dtype=torch.bool)
txt_fea = torch.randn(batch_size, text_len, t5_feat_dim, device=device, dtype=dtype)
with torch.no_grad():
output = model(x, t, y, mask, txt_fea)
expected_shape = (batch_size, num_frames, in_channels, height, width)
assert output.shape == expected_shape
assert output.dtype == dtype
assert output.device == device
def test_model_parameter_count_1_3b(self, sample_config_1_3b):
"""Test parameter count is reasonable for 1.3B model."""
with torch.device('meta'):
model = WanModel(**sample_config_1_3b)
total_params = sum(p.numel() for p in model.parameters())
# Should be around 1.3B parameters (allow some variance)
assert 1.0e9 < total_params < 2.0e9, f"Expected ~1.3B params, got {total_params:,}"
def test_model_parameter_count_14b(self, sample_config_14b):
"""Test parameter count is reasonable for 14B model."""
with torch.device('meta'):
model = WanModel(**sample_config_14b)
total_params = sum(p.numel() for p in model.parameters())
# Should be around 14B parameters (allow some variance)
assert 10e9 < total_params < 20e9, f"Expected ~14B params, got {total_params:,}"
def test_model_no_nan_output(self, sample_config_1_3b, device, dtype):
"""Test that model output doesn't contain NaN values."""
config = sample_config_1_3b.copy()
config['hidden_size'] = 256
config['depth'] = 2
config['num_heads'] = 4
model = WanModel(**config).to(device).to(dtype)
model.eval()
batch_size = 1
num_frames = 4
height = 16
width = 16
in_channels = config['in_channels']
text_len = config['text_len']
t5_feat_dim = config['t5_feat_dim']
cap_feat_dim = config['cap_feat_dim']
x = torch.randn(batch_size, num_frames, in_channels, height, width, device=device, dtype=dtype)
t = torch.randn(batch_size, device=device, dtype=dtype)
y = torch.randn(batch_size, 1, cap_feat_dim, device=device, dtype=dtype)
mask = torch.ones(batch_size, text_len, device=device, dtype=torch.bool)
txt_fea = torch.randn(batch_size, text_len, t5_feat_dim, device=device, dtype=dtype)
with torch.no_grad():
output = model(x, t, y, mask, txt_fea)
assert not torch.isnan(output).any()
assert not torch.isinf(output).any()
def test_model_eval_mode(self, sample_config_1_3b, device):
"""Test that model can be set to eval mode."""
with torch.device('meta'):
model = WanModel(**sample_config_1_3b)
model.eval()
assert not model.training
def test_model_train_mode(self, sample_config_1_3b, device):
"""Test that model can be set to train mode."""
with torch.device('meta'):
model = WanModel(**sample_config_1_3b)
model.train()
assert model.training
def test_model_config_attributes(self, sample_config_1_3b):
"""Test that model has correct configuration attributes."""
with torch.device('meta'):
model = WanModel(**sample_config_1_3b)
assert hasattr(model, 'patch_size')
assert hasattr(model, 'in_channels')
assert hasattr(model, 'hidden_size')
assert hasattr(model, 'depth')
assert hasattr(model, 'num_heads')
assert model.patch_size == sample_config_1_3b['patch_size']
assert model.in_channels == sample_config_1_3b['in_channels']
@pytest.mark.parametrize("batch_size", [1, 2, 4])
def test_model_various_batch_sizes(self, sample_config_1_3b, device, dtype, batch_size):
"""Test model with various batch sizes."""
config = sample_config_1_3b.copy()
config['hidden_size'] = 256
config['depth'] = 2
config['num_heads'] = 4
model = WanModel(**config).to(device).to(dtype)
model.eval()
num_frames = 4
height = 16
width = 16
in_channels = config['in_channels']
text_len = config['text_len']
t5_feat_dim = config['t5_feat_dim']
cap_feat_dim = config['cap_feat_dim']
x = torch.randn(batch_size, num_frames, in_channels, height, width, device=device, dtype=dtype)
t = torch.randn(batch_size, device=device, dtype=dtype)
y = torch.randn(batch_size, 1, cap_feat_dim, device=device, dtype=dtype)
mask = torch.ones(batch_size, text_len, device=device, dtype=torch.bool)
txt_fea = torch.randn(batch_size, text_len, t5_feat_dim, device=device, dtype=dtype)
with torch.no_grad():
output = model(x, t, y, mask, txt_fea)
assert output.shape[0] == batch_size

153
tests/test_pipelines.py Normal file
View File

@ -0,0 +1,153 @@
"""
Integration tests for Wan2.1 pipelines (T2V, I2V, FLF2V, VACE).
Copyright (c) 2025 Kuaishou. All rights reserved.
Note: These tests require model checkpoints and are marked as integration tests.
Run with: pytest -m integration
"""
import pytest
import torch
@pytest.mark.integration
@pytest.mark.requires_model
class TestText2VideoPipeline:
"""Integration tests for Text-to-Video pipeline."""
def test_t2v_pipeline_imports(self):
"""Test that T2V pipeline can be imported."""
from wan.text2video import WanT2V
assert WanT2V is not None
def test_t2v_pipeline_initialization(self):
"""Test T2V pipeline initialization (meta device, no weights)."""
from wan.text2video import WanT2V
# This tests the interface without loading actual weights
# Real tests would require model checkpoints
assert callable(WanT2V)
@pytest.mark.integration
@pytest.mark.requires_model
class TestImage2VideoPipeline:
"""Integration tests for Image-to-Video pipeline."""
def test_i2v_pipeline_imports(self):
"""Test that I2V pipeline can be imported."""
from wan.image2video import WanI2V
assert WanI2V is not None
def test_i2v_pipeline_initialization(self):
"""Test I2V pipeline initialization (meta device, no weights)."""
from wan.image2video import WanI2V
assert callable(WanI2V)
@pytest.mark.integration
@pytest.mark.requires_model
class TestFirstLastFrame2VideoPipeline:
"""Integration tests for First-Last-Frame-to-Video pipeline."""
def test_flf2v_pipeline_imports(self):
"""Test that FLF2V pipeline can be imported."""
from wan.first_last_frame2video import WanFLF2V
assert WanFLF2V is not None
def test_flf2v_pipeline_initialization(self):
"""Test FLF2V pipeline initialization (meta device, no weights)."""
from wan.first_last_frame2video import WanFLF2V
assert callable(WanFLF2V)
@pytest.mark.integration
@pytest.mark.requires_model
class TestVACEPipeline:
"""Integration tests for VACE (Video Creation & Editing) pipeline."""
def test_vace_pipeline_imports(self):
"""Test that VACE pipeline can be imported."""
from wan.vace import WanVace
assert WanVace is not None
def test_vace_pipeline_initialization(self):
"""Test VACE pipeline initialization (meta device, no weights)."""
from wan.vace import WanVace
assert callable(WanVace)
class TestPipelineConfigs:
"""Test pipeline configuration loading."""
def test_t2v_14b_config_loads(self):
"""Test that T2V 14B config can be loaded."""
from wan.configs.t2v_14B import get_config
config = get_config()
assert config is not None
assert 'hidden_size' in config
assert config['hidden_size'] == 3072
def test_t2v_1_3b_config_loads(self):
"""Test that T2V 1.3B config can be loaded."""
from wan.configs.t2v_1_3B import get_config
config = get_config()
assert config is not None
assert 'hidden_size' in config
assert config['hidden_size'] == 1536
def test_i2v_14b_config_loads(self):
"""Test that I2V 14B config can be loaded."""
from wan.configs.i2v_14B import get_config
config = get_config()
assert config is not None
assert 'hidden_size' in config
def test_i2v_1_3b_config_loads(self):
"""Test that I2V 1.3B config can be loaded."""
from wan.configs.i2v_1_3B import get_config
config = get_config()
assert config is not None
assert 'hidden_size' in config
def test_all_configs_have_required_keys(self):
"""Test that all configs have required keys."""
from wan.configs.t2v_14B import get_config as get_t2v_14b
from wan.configs.t2v_1_3B import get_config as get_t2v_1_3b
from wan.configs.i2v_14B import get_config as get_i2v_14b
from wan.configs.i2v_1_3B import get_config as get_i2v_1_3b
required_keys = [
'patch_size', 'in_channels', 'hidden_size', 'depth',
'num_heads', 'mlp_ratio', 'learn_sigma'
]
for config_fn in [get_t2v_14b, get_t2v_1_3b, get_i2v_14b, get_i2v_1_3b]:
config = config_fn()
for key in required_keys:
assert key in config, f"Missing key {key} in config"
class TestDistributed:
"""Test distributed training utilities."""
def test_fsdp_imports(self):
"""Test that FSDP utilities can be imported."""
from wan.distributed.fsdp import WanFSDP
assert WanFSDP is not None
def test_context_parallel_imports(self):
"""Test that context parallel utilities can be imported."""
try:
from wan.distributed.xdit_context_parallel import xFuserWanModelArgs
assert xFuserWanModelArgs is not None
except ImportError:
pytest.skip("xDiT context parallel not available")

190
tests/test_utils.py Normal file
View File

@ -0,0 +1,190 @@
"""
Unit tests for utility functions in Wan2.1.
Copyright (c) 2025 Kuaishou. All rights reserved.
"""
import pytest
import torch
import tempfile
from pathlib import Path
from wan.utils.utils import video_to_torch_cached, image_to_torch_cached
class TestUtilityFunctions:
"""Test suite for utility functions."""
def test_image_to_torch_cached_basic(self, temp_dir):
"""Test basic image loading and caching."""
# Create a dummy image file using PIL
try:
from PIL import Image
import numpy as np
# Create a simple test image
img_array = np.random.randint(0, 255, (64, 64, 3), dtype=np.uint8)
img = Image.fromarray(img_array)
img_path = temp_dir / "test_image.png"
img.save(img_path)
# Load image with caching
tensor = image_to_torch_cached(str(img_path))
assert isinstance(tensor, torch.Tensor)
assert tensor.ndim == 3 # CHW format
assert tensor.shape[0] == 3 # RGB channels
assert tensor.dtype == torch.float32
assert tensor.min() >= 0.0
assert tensor.max() <= 1.0
except ImportError:
pytest.skip("PIL not available")
def test_image_to_torch_cached_resize(self, temp_dir):
"""Test image loading with resizing."""
try:
from PIL import Image
import numpy as np
# Create a test image
img_array = np.random.randint(0, 255, (128, 128, 3), dtype=np.uint8)
img = Image.fromarray(img_array)
img_path = temp_dir / "test_image.png"
img.save(img_path)
# Load and resize
target_size = (64, 64)
tensor = image_to_torch_cached(str(img_path), size=target_size)
assert tensor.shape[1] == target_size[0] # height
assert tensor.shape[2] == target_size[1] # width
except ImportError:
pytest.skip("PIL not available")
def test_image_to_torch_nonexistent_file(self):
"""Test that loading a nonexistent image raises an error."""
with pytest.raises(Exception):
image_to_torch_cached("/nonexistent/path/image.png")
def test_video_to_torch_cached_basic(self, temp_dir):
"""Test basic video loading (if av is available)."""
try:
import av
import numpy as np
# Create a simple test video
video_path = temp_dir / "test_video.mp4"
container = av.open(str(video_path), mode='w')
stream = container.add_stream('mpeg4', rate=30)
stream.width = 64
stream.height = 64
stream.pix_fmt = 'yuv420p'
for i in range(10):
frame = av.VideoFrame(64, 64, 'rgb24')
frame_array = np.random.randint(0, 255, (64, 64, 3), dtype=np.uint8)
frame.planes[0].update(frame_array)
packet = stream.encode(frame)
if packet:
container.mux(packet)
# Flush remaining packets
for packet in stream.encode():
container.mux(packet)
container.close()
# Load video with caching
tensor = video_to_torch_cached(str(video_path))
assert isinstance(tensor, torch.Tensor)
assert tensor.ndim == 4 # TCHW format
assert tensor.shape[1] == 3 # RGB channels
assert tensor.dtype == torch.float32
except ImportError:
pytest.skip("av library not available")
def test_video_to_torch_nonexistent_file(self):
"""Test that loading a nonexistent video raises an error."""
with pytest.raises(Exception):
video_to_torch_cached("/nonexistent/path/video.mp4")
class TestPromptExtension:
"""Test suite for prompt extension utilities."""
def test_prompt_extend_imports(self):
"""Test that prompt extension modules can be imported."""
try:
from wan.utils.prompt_extend import extend_prompt_with_qwen, extend_prompt_with_dashscope
assert extend_prompt_with_qwen is not None
assert extend_prompt_with_dashscope is not None
except ImportError as e:
pytest.fail(f"Failed to import prompt extension: {e}")
def test_prompt_extend_qwen_basic(self):
"""Test basic Qwen prompt extension (without model)."""
try:
from wan.utils.prompt_extend import extend_prompt_with_qwen
# This will likely fail without a model, but we're testing the interface
# In a real test, you'd mock the model
assert callable(extend_prompt_with_qwen)
except ImportError:
pytest.skip("Prompt extension not available")
class TestFMSolvers:
"""Test suite for flow matching solvers."""
def test_fm_solver_imports(self):
"""Test that FM solver modules can be imported."""
from wan.utils.fm_solvers import FlowMatchingDPMSolver
from wan.utils.fm_solvers_unipc import FlowMatchingUniPCSolver
assert FlowMatchingDPMSolver is not None
assert FlowMatchingUniPCSolver is not None
def test_dpm_solver_initialization(self):
"""Test DPM solver initialization."""
from wan.utils.fm_solvers import FlowMatchingDPMSolver
solver = FlowMatchingDPMSolver(
num_steps=20,
order=2,
skip_type='time_uniform',
method='multistep',
)
assert solver is not None
assert solver.num_steps == 20
def test_unipc_solver_initialization(self):
"""Test UniPC solver initialization."""
from wan.utils.fm_solvers_unipc import FlowMatchingUniPCSolver
solver = FlowMatchingUniPCSolver(
num_steps=20,
order=2,
skip_type='time_uniform',
)
assert solver is not None
assert solver.num_steps == 20
def test_solver_get_timesteps(self):
"""Test that solver can generate timesteps."""
from wan.utils.fm_solvers import FlowMatchingDPMSolver
solver = FlowMatchingDPMSolver(
num_steps=10,
order=2,
)
timesteps = solver.get_time_steps()
assert len(timesteps) > 0
assert all(0 <= t <= 1 for t in timesteps)

220
tests/test_vae.py Normal file
View File

@ -0,0 +1,220 @@
"""
Unit tests for WanVAE in Wan2.1.
Copyright (c) 2025 Kuaishou. All rights reserved.
"""
import pytest
import torch
from wan.modules.vae import WanVAE_
class TestWanVAE:
"""Test suite for WanVAE (3D Causal VAE)."""
def test_vae_initialization(self, sample_vae_config):
"""Test VAE initialization."""
with torch.device('meta'):
vae = WanVAE_(**sample_vae_config)
assert vae is not None
assert hasattr(vae, 'encoder')
assert hasattr(vae, 'decoder')
assert vae.temporal_compress_level == sample_vae_config['temporal_compress_level']
def test_vae_encode_shape(self, sample_vae_config, device, dtype):
"""Test VAE encoding produces correct output shape."""
# Use smaller config for faster testing
config = sample_vae_config.copy()
config['encoder_config']['ch'] = 32
config['encoder_config']['ch_mult'] = [1, 2]
config['encoder_config']['num_res_blocks'] = 1
config['decoder_config']['ch'] = 32
config['decoder_config']['ch_mult'] = [1, 2]
config['decoder_config']['num_res_blocks'] = 1
vae = WanVAE_(**config).to(device).to(dtype)
vae.eval()
batch_size = 1
channels = 3
num_frames = 8
height = 64
width = 64
x = torch.randn(batch_size, channels, num_frames, height, width, device=device, dtype=dtype)
with torch.no_grad():
encoded = vae.encode(x)
# Check output shape after encoding
z_channels = config['encoder_config']['z_channels']
temporal_compress = config['temporal_compress_level']
spatial_compress = 2 ** (len(config['encoder_config']['ch_mult']) - 1)
expected_t = num_frames // temporal_compress
expected_h = height // spatial_compress
expected_w = width // spatial_compress
assert encoded.shape == (batch_size, z_channels, expected_t, expected_h, expected_w)
def test_vae_decode_shape(self, sample_vae_config, device, dtype):
"""Test VAE decoding produces correct output shape."""
config = sample_vae_config.copy()
config['encoder_config']['ch'] = 32
config['encoder_config']['ch_mult'] = [1, 2]
config['encoder_config']['num_res_blocks'] = 1
config['decoder_config']['ch'] = 32
config['decoder_config']['ch_mult'] = [1, 2]
config['decoder_config']['num_res_blocks'] = 1
vae = WanVAE_(**config).to(device).to(dtype)
vae.eval()
batch_size = 1
z_channels = config['encoder_config']['z_channels']
num_frames = 2
height = 32
width = 32
z = torch.randn(batch_size, z_channels, num_frames, height, width, device=device, dtype=dtype)
with torch.no_grad():
decoded = vae.decode(z)
# Check output shape after decoding
out_channels = config['decoder_config']['out_ch']
temporal_compress = config['temporal_compress_level']
spatial_compress = 2 ** (len(config['decoder_config']['ch_mult']) - 1)
expected_t = num_frames * temporal_compress
expected_h = height * spatial_compress
expected_w = width * spatial_compress
assert decoded.shape == (batch_size, out_channels, expected_t, expected_h, expected_w)
def test_vae_encode_decode_consistency(self, sample_vae_config, device, dtype):
"""Test that encode then decode produces similar output."""
config = sample_vae_config.copy()
config['encoder_config']['ch'] = 32
config['encoder_config']['ch_mult'] = [1, 2]
config['encoder_config']['num_res_blocks'] = 1
config['decoder_config']['ch'] = 32
config['decoder_config']['ch_mult'] = [1, 2]
config['decoder_config']['num_res_blocks'] = 1
vae = WanVAE_(**config).to(device).to(dtype)
vae.eval()
batch_size = 1
channels = 3
num_frames = 8
height = 64
width = 64
x = torch.randn(batch_size, channels, num_frames, height, width, device=device, dtype=dtype)
with torch.no_grad():
encoded = vae.encode(x)
decoded = vae.decode(encoded)
# Decoded output should have same shape as input
assert decoded.shape == x.shape
def test_vae_no_nan_encode(self, sample_vae_config, device, dtype):
"""Test that VAE encoding doesn't produce NaN values."""
config = sample_vae_config.copy()
config['encoder_config']['ch'] = 32
config['encoder_config']['ch_mult'] = [1, 2]
config['encoder_config']['num_res_blocks'] = 1
config['decoder_config']['ch'] = 32
config['decoder_config']['ch_mult'] = [1, 2]
config['decoder_config']['num_res_blocks'] = 1
vae = WanVAE_(**config).to(device).to(dtype)
vae.eval()
batch_size = 1
channels = 3
num_frames = 8
height = 64
width = 64
x = torch.randn(batch_size, channels, num_frames, height, width, device=device, dtype=dtype)
with torch.no_grad():
encoded = vae.encode(x)
assert not torch.isnan(encoded).any()
assert not torch.isinf(encoded).any()
def test_vae_no_nan_decode(self, sample_vae_config, device, dtype):
"""Test that VAE decoding doesn't produce NaN values."""
config = sample_vae_config.copy()
config['encoder_config']['ch'] = 32
config['encoder_config']['ch_mult'] = [1, 2]
config['encoder_config']['num_res_blocks'] = 1
config['decoder_config']['ch'] = 32
config['decoder_config']['ch_mult'] = [1, 2]
config['decoder_config']['num_res_blocks'] = 1
vae = WanVAE_(**config).to(device).to(dtype)
vae.eval()
batch_size = 1
z_channels = config['encoder_config']['z_channels']
num_frames = 2
height = 32
width = 32
z = torch.randn(batch_size, z_channels, num_frames, height, width, device=device, dtype=dtype)
with torch.no_grad():
decoded = vae.decode(z)
assert not torch.isnan(decoded).any()
assert not torch.isinf(decoded).any()
@pytest.mark.parametrize("num_frames", [4, 8, 16])
def test_vae_various_frame_counts(self, sample_vae_config, device, dtype, num_frames):
"""Test VAE with various frame counts."""
config = sample_vae_config.copy()
config['encoder_config']['ch'] = 32
config['encoder_config']['ch_mult'] = [1, 2]
config['encoder_config']['num_res_blocks'] = 1
config['decoder_config']['ch'] = 32
config['decoder_config']['ch_mult'] = [1, 2]
config['decoder_config']['num_res_blocks'] = 1
vae = WanVAE_(**config).to(device).to(dtype)
vae.eval()
batch_size = 1
channels = 3
height = 64
width = 64
x = torch.randn(batch_size, channels, num_frames, height, width, device=device, dtype=dtype)
with torch.no_grad():
encoded = vae.encode(x)
decoded = vae.decode(encoded)
assert decoded.shape == x.shape
assert not torch.isnan(decoded).any()
def test_vae_eval_mode(self, sample_vae_config):
"""Test that VAE can be set to eval mode."""
with torch.device('meta'):
vae = WanVAE_(**sample_vae_config)
vae.eval()
assert not vae.training
def test_vae_config_attributes(self, sample_vae_config):
"""Test that VAE has correct configuration attributes."""
with torch.device('meta'):
vae = WanVAE_(**sample_vae_config)
assert hasattr(vae, 'temporal_compress_level')
assert vae.temporal_compress_level == sample_vae_config['temporal_compress_level']

2
wan/modules/clip.py
View File

@ -516,7 +516,7 @@ class CLIPModel:
self.model = self.model.eval().requires_grad_(False)
logging.info(f'loading {checkpoint_path}')
self.model.load_state_dict(
torch.load(checkpoint_path, map_location='cpu'))
torch.load(checkpoint_path, map_location='cpu', weights_only=True))
# init tokenizer
self.tokenizer = HuggingfaceTokenizer(

2
wan/modules/t5.py
View File

@ -493,7 +493,7 @@ class T5EncoderModel:
dtype=dtype,
device=device).eval().requires_grad_(False)
logging.info(f'loading {checkpoint_path}')
model.load_state_dict(torch.load(checkpoint_path, map_location='cpu'))
model.load_state_dict(torch.load(checkpoint_path, map_location='cpu', weights_only=True))
self.model = model
if shard_fn is not None:
self.model = shard_fn(self.model, sync_module_states=False)

2
wan/modules/vae.py
View File

@ -611,7 +611,7 @@ def _video_vae(pretrained_path=None, z_dim=None, device='cpu', **kwargs):
# load checkpoint
logging.info(f'loading {pretrained_path}')
model.load_state_dict(
torch.load(pretrained_path, map_location=device), assign=True)
torch.load(pretrained_path, map_location=device, weights_only=True), assign=True)
return model