From 15553f770699df214efc4131817ba808e9ba0d74 Mon Sep 17 00:00:00 2001 From: William Lin Date: Tue, 8 Apr 2025 11:38:47 -0700 Subject: [PATCH] [CI] Use pre-commit to run linter (#321) Co-authored-by: kevin314 --- .github/scripts/runpod_api.py | 11 +- .github/scripts/runpod_cleanup.py | 5 +- .github/workflows/codespell.yml | 45 --- .github/workflows/fastvideo-publish.yml | 10 +- .github/workflows/matchers/actionlint.json | 17 ++ .github/workflows/matchers/mypy.json | 16 + .github/workflows/pr-test.yml | 12 +- .github/workflows/pre-commit.yml | 18 ++ .github/workflows/ruff.yml | 50 ---- .github/workflows/sta-publish.yml | 8 +- .github/workflows/test.yml | 4 +- .github/workflows/yapf.yml | 38 --- .pre-commit-config.yaml | 80 +++++ README.md | 50 +++- docs/data_preprocess.md | 13 +- fastvideo/v1/attention/backends/sdpa.py | 5 +- .../device_communicators/cuda_communicator.py | 3 +- .../v1/layers/vocab_parallel_embedding.py | 16 +- fastvideo/v1/models/dits/hunyuanvideo.py | 8 +- fastvideo/v1/models/encoders/clip.py | 8 +- fastvideo/v1/models/loader/fsdp_load.py | 2 +- fastvideo/v1/models/registry.py | 3 +- fastvideo/v1/models/vaes/hunyuanvae.py | 7 +- fastvideo/v1/pipelines/README.md | 2 +- fastvideo/v1/tests/conftest.py | 20 +- .../v1/tests/encoders/test_clip_encoder.py | 74 +++-- .../v1/tests/encoders/test_llama_encoder.py | 89 +++--- .../v1/tests/old_tests/test_hunyuan_vae.py | 16 +- .../v1/tests/old_tests/test_hunyuanvideo.py | 33 ++- .../tests/old_tests/test_hunyuanvideo_load.py | 47 ++- .../v1/tests/old_tests/test_t5_encoder.py | 127 ++++---- fastvideo/v1/tests/old_tests/test_wan_vae.py | 102 ++++--- fastvideo/v1/tests/old_tests/tp_example.py | 18 +- fastvideo/v1/tests/ssim/README.md | 2 +- fastvideo/v1/tests/ssim/compute_ssim.py | 11 +- .../tests/ssim/test_inference_similarity.py | 13 +- fastvideo/v1/utils.py | 8 +- format.sh | 275 ------------------ pyproject.toml | 15 +- 39 files changed, 525 insertions(+), 756 deletions(-) delete mode 100644 .github/workflows/codespell.yml create mode 100644 .github/workflows/matchers/actionlint.json create mode 100644 .github/workflows/matchers/mypy.json create mode 100644 .github/workflows/pre-commit.yml delete mode 100644 .github/workflows/ruff.yml delete mode 100644 .github/workflows/yapf.yml create mode 100644 .pre-commit-config.yaml delete mode 100644 format.sh diff --git a/.github/scripts/runpod_api.py b/.github/scripts/runpod_api.py index 269bbe677..9836caaba 100644 --- a/.github/scripts/runpod_api.py +++ b/.github/scripts/runpod_api.py @@ -1,10 +1,11 @@ -import os -import json -import time -import requests -import sys import argparse +import json +import os import subprocess +import sys +import time + +import requests def parse_arguments(): diff --git a/.github/scripts/runpod_cleanup.py b/.github/scripts/runpod_cleanup.py index df4041bd3..eef991af7 100644 --- a/.github/scripts/runpod_cleanup.py +++ b/.github/scripts/runpod_cleanup.py @@ -1,9 +1,10 @@ import json import os import sys -import requests import uuid +import requests + API_KEY = os.environ['RUNPOD_API_KEY'] RUN_ID = os.environ.get('GITHUB_RUN_ID', str(uuid.uuid4())) PODS_API = "https://rest.runpod.io/v1/pods" @@ -19,7 +20,7 @@ def get_job_ids(): try: job_ids = json.loads(job_ids_str) if not isinstance(job_ids, list): - print(f"Error: JOB_IDS is not a list.") + print("Error: JOB_IDS is not a list.") sys.exit(1) return job_ids except json.JSONDecodeError as e: diff --git a/.github/workflows/codespell.yml b/.github/workflows/codespell.yml deleted file mode 100644 index 06c2ff9e2..000000000 --- a/.github/workflows/codespell.yml +++ /dev/null @@ -1,45 +0,0 @@ -name: codespell - -on: - # Trigger the workflow on push or pull request, - # but only for the main branch - push: - branches: - - main - paths: - - "**/*.py" - - "**/*.md" - - "**/*.rst" - - pyproject.toml - - requirements-lint.txt - - .github/workflows/codespell.yml - pull_request: - branches: - - main - paths: - - "**/*.py" - - "**/*.md" - - "**/*.rst" - - pyproject.toml - - requirements-lint.txt - - .github/workflows/codespell.yml - -jobs: - codespell: - runs-on: ubuntu-latest - steps: - - name: Check out repository - uses: actions/checkout@v3 - - - name: Set up Python - uses: actions/setup-python@v4 - with: - python-version: '3.12' # or any version you need - - name: Install dependencies - run: | - python -m pip install --upgrade pip - pip install -e ".[lint]" - - name: Spelling check with codespell - run: | - # Refer to the above environment variable here - codespell --toml pyproject.toml $CODESPELL_EXCLUDES diff --git a/.github/workflows/fastvideo-publish.yml b/.github/workflows/fastvideo-publish.yml index f7e9391d9..bb1f57441 100644 --- a/.github/workflows/fastvideo-publish.yml +++ b/.github/workflows/fastvideo-publish.yml @@ -15,7 +15,7 @@ jobs: new-version: ${{ steps.check-version.outputs.new-version }} steps: - name: Checkout code - uses: actions/checkout@v3 + uses: actions/checkout@v4 with: fetch-depth: 2 @@ -23,11 +23,11 @@ jobs: id: check-version run: | # Get current commit's version - NEW_VERSION=$(grep -oP 'version\s*=\s*"\K[^"]+' pyproject.toml) + NEW_VERSION=$(grep -oP "version\\s*=\\s*\"\\K[^\"]+\"" pyproject.toml) echo "New version: $NEW_VERSION" # Get previous version from git history - OLD_VERSION=$(git show HEAD~1:./pyproject.toml | grep -oP 'version\s*=\s*"\K[^"]+' || echo "0.0.0") + OLD_VERSION=$(git show HEAD~1:./pyproject.toml | grep -oP "version\\s*=\\s*\"\\K[^\"]+\"" || echo "0.0.0") echo "Old version: $OLD_VERSION" if [ "$NEW_VERSION" != "$OLD_VERSION" ]; then @@ -48,10 +48,10 @@ jobs: steps: - name: Checkout code - uses: actions/checkout@v3 + uses: actions/checkout@v4 - name: Set up Python - uses: actions/setup-python@v4 + uses: actions/setup-python@v5 with: python-version: '3.10' diff --git a/.github/workflows/matchers/actionlint.json b/.github/workflows/matchers/actionlint.json new file mode 100644 index 000000000..4613e1617 --- /dev/null +++ b/.github/workflows/matchers/actionlint.json @@ -0,0 +1,17 @@ +{ + "problemMatcher": [ + { + "owner": "actionlint", + "pattern": [ + { + "regexp": "^(?:\\x1b\\[\\d+m)?(.+?)(?:\\x1b\\[\\d+m)*:(?:\\x1b\\[\\d+m)*(\\d+)(?:\\x1b\\[\\d+m)*:(?:\\x1b\\[\\d+m)*(\\d+)(?:\\x1b\\[\\d+m)*: (?:\\x1b\\[\\d+m)*(.+?)(?:\\x1b\\[\\d+m)* \\[(.+?)\\]$", + "file": 1, + "line": 2, + "column": 3, + "message": 4, + "code": 5 + } + ] + } + ] +} diff --git a/.github/workflows/matchers/mypy.json b/.github/workflows/matchers/mypy.json new file mode 100644 index 000000000..f048fce52 --- /dev/null +++ b/.github/workflows/matchers/mypy.json @@ -0,0 +1,16 @@ +{ + "problemMatcher": [ + { + "owner": "mypy", + "pattern": [ + { + "regexp": "^(.+):(\\d+):\\s(error|warning):\\s(.+)$", + "file": 1, + "line": 2, + "severity": 3, + "message": 4 + } + ] + } + ] +} diff --git a/.github/workflows/pr-test.yml b/.github/workflows/pr-test.yml index 802a4395f..7928f2865 100644 --- a/.github/workflows/pr-test.yml +++ b/.github/workflows/pr-test.yml @@ -30,14 +30,18 @@ concurrency: cancel-in-progress: true jobs: + pre-commit: + uses: ./.github/workflows/pre-commit.yml + change-filter: runs-on: ubuntu-latest + needs: pre-commit if: ${{ github.event.pull_request.draft == false || github.event_name == 'workflow_dispatch' }} outputs: encoder-test: ${{ steps.filter.outputs.encoder-test }} steps: - uses: actions/checkout@v4 - - uses: dorny/paths-filter@v2 + - uses: dorny/paths-filter@v3 id: filter with: filters: | @@ -58,7 +62,7 @@ jobs: uses: actions/checkout@v4 - name: Set up Python - uses: actions/setup-python@v4 + uses: actions/setup-python@v5 with: python-version: "3.10" @@ -107,7 +111,7 @@ jobs: uses: actions/checkout@v4 - name: Set up Python - uses: actions/setup-python@v4 + uses: actions/setup-python@v5 with: python-version: "3.10" @@ -154,7 +158,7 @@ jobs: uses: actions/checkout@v4 - name: Set up Python - uses: actions/setup-python@v4 + uses: actions/setup-python@v5 with: python-version: "3.10" diff --git a/.github/workflows/pre-commit.yml b/.github/workflows/pre-commit.yml new file mode 100644 index 000000000..964cacbb2 --- /dev/null +++ b/.github/workflows/pre-commit.yml @@ -0,0 +1,18 @@ +name: pre-commit + +on: + workflow_call: + +jobs: + pre-commit: + runs-on: ubuntu-latest + steps: + - uses: actions/checkout@v4 + - uses: actions/setup-python@v5 + with: + python-version: "3.10" + - run: echo "::add-matcher::.github/workflows/matchers/actionlint.json" + - run: echo "::add-matcher::.github/workflows/matchers/mypy.json" + - uses: pre-commit/action@v3.0.1 + with: + extra_args: --all-files --hook-stage manual \ No newline at end of file diff --git a/.github/workflows/ruff.yml b/.github/workflows/ruff.yml deleted file mode 100644 index 954247262..000000000 --- a/.github/workflows/ruff.yml +++ /dev/null @@ -1,50 +0,0 @@ -name: ruff - -on: - # Trigger the workflow on push or pull request, - # but only for the main branch - push: - branches: - - main - paths: - - "**/*.py" - - pyproject.toml - - requirements-lint.txt - - .github/workflows/matchers/ruff.json - - .github/workflows/ruff.yml - pull_request: - branches: - - main - # This workflow is only relevant when one of the following files changes. - # However, we have github configured to expect and require this workflow - # to run and pass before github with auto-merge a pull request. Until github - # allows more flexible auto-merge policy, we can just run this on every PR. - # It doesn't take that long to run, anyway. - #paths: - # - "**/*.py" - # - pyproject.toml - # - requirements-lint.txt - # - .github/workflows/matchers/ruff.json - # - .github/workflows/ruff.yml - -jobs: - ruff: - runs-on: ubuntu-latest - steps: - - name: Check out repository - uses: actions/checkout@v3 - - - name: Set up Python - uses: actions/setup-python@v4 - with: - python-version: '3.12' # or any version you need - - name: Install dependencies - run: | - python -m pip install --upgrade pip - pip install -e ".[lint]" - - name: Analysing the code with ruff - run: | - ruff check . - - name: Run isort - run: | - isort . --check-only diff --git a/.github/workflows/sta-publish.yml b/.github/workflows/sta-publish.yml index cdc4e412a..660f1a0ce 100644 --- a/.github/workflows/sta-publish.yml +++ b/.github/workflows/sta-publish.yml @@ -15,7 +15,7 @@ jobs: new-version: ${{ steps.check-version.outputs.new-version }} steps: - name: Checkout code - uses: actions/checkout@v3 + uses: actions/checkout@v4 with: fetch-depth: 2 @@ -43,7 +43,7 @@ jobs: build_wheels: name: Build Wheel needs: check-version-change - if: needs.check-version-change.outputs.version-changed == 'true' + if: ${{ needs.check-version-change.outputs.version-changed == 'true' }} runs-on: ${{ matrix.os }} strategy: @@ -144,8 +144,8 @@ jobs: publish_package: name: Publish package - needs: [build_wheels] - if: needs.check-version-change.outputs.version-changed == 'true' + needs: [build_wheels, check-version-change] + if: ${{ needs.check-version-change.outputs.version-changed == 'true' }} runs-on: ubuntu-22.04 permissions: id-token: write # Needed for OIDC Trusted Publishing diff --git a/.github/workflows/test.yml b/.github/workflows/test.yml index 10065edc3..a8b3ddf9a 100644 --- a/.github/workflows/test.yml +++ b/.github/workflows/test.yml @@ -11,10 +11,10 @@ jobs: runs-on: ubuntu-latest steps: - name: Check out repository - uses: actions/checkout@v3 + uses: actions/checkout@v4 - name: Set up Python - uses: actions/setup-python@v4 + uses: actions/setup-python@v5 with: python-version: '3.12' # or any version you need diff --git a/.github/workflows/yapf.yml b/.github/workflows/yapf.yml deleted file mode 100644 index 3908e9522..000000000 --- a/.github/workflows/yapf.yml +++ /dev/null @@ -1,38 +0,0 @@ -name: yapf - -on: - # Trigger the workflow on push or pull request, - # but only for the main branch - push: - branches: - - main - paths: - - "**/*.py" - - .github/workflows/yapf.yml - pull_request: - branches: - - main - paths: - - "**/*.py" - - .github/workflows/yapf.yml - -jobs: - yapf: - runs-on: ubuntu-latest - steps: - - name: Check out repository - uses: actions/checkout@v3 - - - name: Set up Python - uses: actions/setup-python@v4 - with: - python-version: '3.12' # or any version you need - - - name: Install dependencies - run: | - python -m pip install --upgrade pip - pip install yapf==0.32.0 - pip install toml==0.10.2 - - name: Running yapf - run: | - yapf --diff --recursive . diff --git a/.pre-commit-config.yaml b/.pre-commit-config.yaml new file mode 100644 index 000000000..310316204 --- /dev/null +++ b/.pre-commit-config.yaml @@ -0,0 +1,80 @@ +default_stages: + - pre-commit # Run locally + - manual # Run in CI +exclude: | + (?x)( + fastvideo/v1/third_party/.*| + csrc/.*| + assets/.*| + tests/.*| + demo/.*| + predict\.py| + scripts/.*| + fastvideo/data_preprocess/.*| + fastvideo/dataset/.*| + fastvideo/distill/.*| + fastvideo/distill\.py| + fastvideo/distill_adv\.py| + fastvideo/models/.*| + fastvideo/sample/.*| + fastvideo/train\.py| + fastvideo/utils/.*| + .github/workflows/fastvideo-publish.yml| + .github/workflows/sta-publish.yml + ) +repos: +- repo: https://github.com/google/yapf + rev: v0.43.0 + hooks: + - id: yapf + args: [--in-place, --verbose] + additional_dependencies: [toml] # TODO: Remove when yapf is upgraded +- repo: https://github.com/astral-sh/ruff-pre-commit + rev: v0.11.4 + hooks: + - id: ruff + args: [--output-format, github, --fix] +- repo: https://github.com/codespell-project/codespell + rev: v2.4.1 + hooks: + - id: codespell + additional_dependencies: ['tomli'] + args: ['--toml', 'pyproject.toml'] +- repo: https://github.com/PyCQA/isort + rev: 0a0b7a830386ba6a31c2ec8316849ae4d1b8240d # 6.0.0 + hooks: + - id: isort +- repo: https://github.com/jackdewinter/pymarkdown + rev: v0.9.29 + hooks: + - id: pymarkdown + args: [fix] +- repo: https://github.com/rhysd/actionlint + rev: v1.7.7 + hooks: + - id: actionlint +- repo: https://github.com/pre-commit/mirrors-mypy + rev: v1.15.0 + hooks: + - id: mypy + args: [--python-version, '3.10', --follow-imports, "skip", ] + additional_dependencies: [types-cachetools, types-setuptools, types-PyYAML, types-requests] +- repo: local + hooks: + - id: check-filenames + name: Check for spaces in all filenames + entry: bash + args: + - -c + - 'git ls-files | grep -v "^fastvideo/v1/tests/ssim/reference_videos/" | grep " " && echo "Filenames should not contain spaces!" && exit 1 || exit 0' + language: system + always_run: true + pass_filenames: false + # Keep `suggestion` last + - id: suggestion + name: Suggestion + entry: bash -c 'echo "To bypass pre-commit hooks, add --no-verify to git commit."' + language: system + verbose: true + pass_filenames: false + # Insert new entries above the `suggestion` entry diff --git a/README.md b/README.md index 257bffe9a..df761146e 100644 --- a/README.md +++ b/README.md @@ -4,19 +4,12 @@ FastVideo is a lightweight framework for accelerating large video diffusion models. -

πŸ€— FastHunyuan | πŸ€— FastMochi | πŸŸ£πŸ’¬ Slack

- - - - https://github.com/user-attachments/assets/79af5fb8-707c-4263-b153-9ab2a01d3ac1 - - FastVideo currently offers: (with more to come) - [NEW!] [Sliding Tile Attention](https://hao-ai-lab.github.io/blogs/sta/). @@ -28,8 +21,6 @@ FastVideo currently offers: (with more to come) Dev in progress and highly experimental. - - ## Change Log - ```2025/02/20```: FastVideo now supports STA on [StepVideo](https://github.com/stepfun-ai/Step-Video-T2V) with 3.4X speedup! - ```2025/02/18```: Release the inference code and kernel for [Sliding Tile Attention](https://hao-ai-lab.github.io/blogs/sta/). @@ -37,9 +28,9 @@ Dev in progress and highly experimental. - ```2024/12/25```: Enable single 4090 inference for `FastHunyuan`, please rerun the installation steps to update the environment. - ```2024/12/17```: `FastVideo` v1.0 is released. - ## πŸ”§ Installation from source The code is tested on Python 3.10.0, CUDA 12.4 and H100. + ``` # Clone FastVideo git clone https://github.com/hao-ai-lab/FastVideo.git && cd FastVideo @@ -50,15 +41,19 @@ pip install -e . # Install Flash Attention (optional) pip install flash-attn==2.7.0.post2 ``` + To try Sliding Tile Attention (optional), please follow the instruction in [csrc/sliding_tile_attention/README.md](csrc/sliding_tile_attention/README.md) to install STA. ## πŸš€ Inference ### Inference StepVideo with Sliding Tile Attention First, download the model: + ``` python scripts/huggingface/download_hf.py --repo_id=stepfun-ai/stepvideo-t2v --local_dir=data/stepvideo-t2v --repo_type=model ``` + Use the following scripts to run inference for StepVideo. When using STA for inference, the generated videos will have dimensions of 204Γ—768Γ—768 (currently, this is the only supported shape). + ```bash sh scripts/inference/inference_stepvideo_STA.sh # Inference stepvideo with STA sh scripts/inference/inference_stepvideo.sh # Inference original stepvideo @@ -66,26 +61,31 @@ sh scripts/inference/inference_stepvideo.sh # Inference original stepvideo ### Inference HunyuanVideo with Sliding Tile Attention First, download the model: + ```bash python scripts/huggingface/download_hf.py --repo_id=FastVideo/hunyuan --local_dir=data/hunyuan --repo_type=model ``` + We provide two examples in the following script to run inference with STA + [TeaCache](https://github.com/ali-vilab/TeaCache) and STA only. + ```bash sh scripts/inference/inference_hunyuan_STA.sh ``` + ### Video Demos using STA + Teacache Visit our [demo website](https://fast-video.github.io/) to explore our complete collection of examples. We shorten a single video generation process from 945s to 317s on H100. ### Inference FastHunyuan on single RTX4090 We now support NF4 and LLM-INT8 quantized inference using BitsAndBytes for FastHunyuan. With NF4 quantization, inference can be performed on a single RTX 4090 GPU, requiring just 20GB of VRAM. + ```bash # Download the model weight python scripts/huggingface/download_hf.py --repo_id=FastVideo/FastHunyuan-diffusers --local_dir=data/FastHunyuan-diffusers --repo_type=model # CLI inference bash scripts/inference/inference_hunyuan_hf_quantization.sh ``` -For more information about the VRAM requirements for BitsAndBytes quantization, please refer to the table below (timing measured on an H100 GPU): +For more information about the VRAM requirements for BitsAndBytes quantization, please refer to the table below (timing measured on an H100 GPU): | Configuration | Memory to Init Transformer | Peak Memory After Init Pipeline (Denoise) | Diffusion Time | End-to-End Time | |--------------------------------|----------------------------|--------------------------------------------|----------------|-----------------| @@ -93,17 +93,17 @@ For more information about the VRAM requirements for BitsAndBytes quantization, | INT8 + Pipeline CPU Offload | 13.911G | 27.979G | 88s | 116.7s | | NF4 + Pipeline CPU Offload | 9.453G | 19.26G | 78s | 114.5s | - - For improved quality in generated videos, we recommend using a GPU with 80GB of memory to run the BF16 model with the original Hunyuan pipeline. To execute the inference, use the following section: ### FastHunyuan + ```bash # Download the model weight python scripts/huggingface/download_hf.py --repo_id=FastVideo/FastHunyuan --local_dir=data/FastHunyuan --repo_type=model # CLI inference bash scripts/inference/inference_hunyuan.sh ``` + You can also inference FastHunyuan in the [official Hunyuan github](https://github.com/Tencent/HunyuanVideo). ### FastMochi @@ -115,79 +115,99 @@ python scripts/huggingface/download_hf.py --repo_id=FastVideo/FastMochi-diffuser bash scripts/inference/inference_mochi_sp.sh ``` - ## 🎯 Distill Our distillation recipe is based on [Phased Consistency Model](https://github.com/G-U-N/Phased-Consistency-Model). We did not find significant improvement using multi-phase distillation, so we keep the one phase setup similar to the original latent consistency model's recipe. We use the [MixKit](https://huggingface.co/datasets/LanguageBind/Open-Sora-Plan-v1.1.0/tree/main/all_mixkit) dataset for distillation. To avoid running the text encoder and VAE during training, we preprocess all data to generate text embeddings and VAE latents. Preprocessing instructions can be found [data_preprocess.md](docs/data_preprocess.md). For convenience, we also provide preprocessed data that can be downloaded directly using the following command: + ```bash python scripts/huggingface/download_hf.py --repo_id=FastVideo/HD-Mixkit-Finetune-Hunyuan --local_dir=data/HD-Mixkit-Finetune-Hunyuan --repo_type=dataset ``` + Next, download the original model weights with: + ```bash python scripts/huggingface/download_hf.py --repo_id=FastVideo/hunyuan --local_dir=data/hunyuan --repo_type=model # original hunyuan python scripts/huggingface/download_hf.py --repo_id=genmo/mochi-1-preview --local_dir=data/mochi --repo_type=model # original mochi ``` + To launch the distillation process, use the following commands: + ``` bash scripts/distill/distill_hunyuan.sh # for hunyuan bash scripts/distill/distill_mochi.sh # for mochi ``` + We also provide an optional script for distillation with adversarial loss, located at `fastvideo/distill_adv.py`. Although we tried adversarial loss, we did not observe significant improvements. ## Finetune ### ⚑ Full Finetune Ensure your data is prepared and preprocessed in the format specified in [data_preprocess.md](docs/data_preprocess.md). For convenience, we also provide a mochi preprocessed Black Myth Wukong data that can be downloaded directly: + ```bash python scripts/huggingface/download_hf.py --repo_id=FastVideo/Mochi-Black-Myth --local_dir=data/Mochi-Black-Myth --repo_type=dataset ``` + Download the original model weights as specified in [Distill Section](#-distill): Then you can run the finetune with: + ``` bash scripts/finetune/finetune_mochi.sh # for mochi ``` + **Note that for finetuning, we did not tune the hyperparameters in the provided script.** ### ⚑ Lora Finetune Hunyuan supports Lora fine-tuning of videos up to 720p. Demos and prompts of Black-Myth-Wukong can be found in [here](https://huggingface.co/FastVideo/Hunyuan-Black-Myth-Wukong-lora-weight). You can download the Lora weight through: + ```bash python scripts/huggingface/download_hf.py --repo_id=FastVideo/Hunyuan-Black-Myth-Wukong-lora-weight --local_dir=data/Hunyuan-Black-Myth-Wukong-lora-weight --repo_type=model ``` + #### Minimum Hardware Requirement - 40 GB GPU memory each for 2 GPUs with lora. - 30 GB GPU memory each for 2 GPUs with CPU offload and lora. - Currently, both Mochi and Hunyuan models support Lora finetuning through diffusers. To generate personalized videos from your own dataset, you'll need to follow three main steps: dataset preparation, finetuning, and inference. #### Dataset Preparation We provide scripts to better help you get started to train on your own characters! You can run this to organize your dataset to get the videos2caption.json before preprocess. Specify your video folder and corresponding caption folder (caption files should be .txt files and have the same name with its video): + ``` python scripts/dataset_preparation/prepare_json_file.py --video_dir data/input_videos/ --prompt_dir data/captions/ --output_path data/output_folder/videos2caption.json --verbose ``` + Also, we provide script to resize your videos: + ``` python scripts/data_preprocess/resize_videos.py ``` + #### Finetuning After basic dataset preparation and preprocess, you can start to finetune your model using Lora: + ``` bash scripts/finetune/finetune_hunyuan_hf_lora.sh ``` + #### Inference For inference with Lora checkpoint, you can run the following scripts with additional parameter `--lora_checkpoint_dir`: + ``` bash scripts/inference/inference_hunyuan_hf.sh ``` + **We also provide scripts for Mochi in the same directory.** #### Finetune with Both Image and Video Our codebase support finetuning with both image and video. + ```bash bash scripts/finetune/finetune_hunyuan.sh bash scripts/finetune/finetune_mochi_lora_mix.sh ``` + For Image-Video Mixture Fine-tuning, make sure to enable the `--group_frame` option in your script. ## πŸ“‘ Development Plan diff --git a/docs/data_preprocess.md b/docs/data_preprocess.md index 710f7de70..5e679a848 100644 --- a/docs/data_preprocess.md +++ b/docs/data_preprocess.md @@ -1,16 +1,16 @@ - - ## 🧱 Data Preprocess To save GPU memory, we precompute text embeddings and VAE latents to eliminate the need to load the text encoder and VAE during training. - We provide a sample dataset to help you get started. Download the source media using the following command: + ```bash python scripts/huggingface/download_hf.py --repo_id=FastVideo/Image-Vid-Finetune-Src --local_dir=data/Image-Vid-Finetune-Src --repo_type=dataset ``` + To preprocess the dataset for fine-tuning or distillation, run: + ``` bash scripts/preprocess/preprocess_mochi_data.sh # for mochi bash scripts/preprocess/preprocess_hunyuan_data.sh # for hunyuan @@ -33,13 +33,16 @@ path_to_dataset_folder/ Format the JSON file as a list, where each item represents a media source: For image media, + ``` { "path": "0.jpg", "cap": ["captions"] } ``` -For video media, + +For video media, + ``` { "path": "1.mp4", @@ -62,7 +65,9 @@ path_to_media_source_foder,path_to_json_file ``` Adjust the `DATA_MERGE_PATH` and `OUTPUT_DIR` in `scripts/preprocess/preprocess_****_data.sh` accordingly and run: + ``` bash scripts/preprocess/preprocess_****_data.sh ``` + The preprocessed data will be put into the `OUTPUT_DIR` and the `videos2caption.json` can be used in finetune and distill scripts. diff --git a/fastvideo/v1/attention/backends/sdpa.py b/fastvideo/v1/attention/backends/sdpa.py index 367964f9d..1230c630c 100644 --- a/fastvideo/v1/attention/backends/sdpa.py +++ b/fastvideo/v1/attention/backends/sdpa.py @@ -67,9 +67,6 @@ class SDPAImpl(AttentionImpl): if query.shape[1] != key.shape[1]: attn_kwargs["enable_gqa"] = True output = torch.nn.functional.scaled_dot_product_attention( - query, - key, - value, - **attn_kwargs) + query, key, value, **attn_kwargs) output = output.transpose(1, 2) return output diff --git a/fastvideo/v1/distributed/device_communicators/cuda_communicator.py b/fastvideo/v1/distributed/device_communicators/cuda_communicator.py index ba50c0dfa..9c0055a35 100644 --- a/fastvideo/v1/distributed/device_communicators/cuda_communicator.py +++ b/fastvideo/v1/distributed/device_communicators/cuda_communicator.py @@ -6,7 +6,8 @@ from typing import Optional import torch from torch.distributed import ProcessGroup -from fastvideo.v1.distributed.device_communicators.base_device_communicator import DeviceCommunicatorBase +from fastvideo.v1.distributed.device_communicators.base_device_communicator import ( + DeviceCommunicatorBase) class CudaCommunicator(DeviceCommunicatorBase): diff --git a/fastvideo/v1/layers/vocab_parallel_embedding.py b/fastvideo/v1/layers/vocab_parallel_embedding.py index 06d6d9a1d..ea5cbcec8 100644 --- a/fastvideo/v1/layers/vocab_parallel_embedding.py +++ b/fastvideo/v1/layers/vocab_parallel_embedding.py @@ -116,17 +116,17 @@ class VocabParallelEmbeddingShardIndices: def __post_init__(self): # sanity checks - assert (self.padded_org_vocab_start_index <= - self.padded_org_vocab_end_index) - assert (self.padded_added_vocab_start_index <= - self.padded_added_vocab_end_index) + assert (self.padded_org_vocab_start_index + <= self.padded_org_vocab_end_index) + assert (self.padded_added_vocab_start_index + <= self.padded_added_vocab_end_index) assert self.org_vocab_start_index <= self.org_vocab_end_index assert self.added_vocab_start_index <= self.added_vocab_end_index assert self.org_vocab_start_index <= self.padded_org_vocab_start_index - assert (self.added_vocab_start_index <= - self.padded_added_vocab_start_index) + assert (self.added_vocab_start_index + <= self.padded_added_vocab_start_index) assert self.org_vocab_end_index <= self.padded_org_vocab_end_index assert self.added_vocab_end_index <= self.padded_added_vocab_end_index @@ -142,8 +142,8 @@ def get_masked_input_and_mask( added_vocab_end_index: int) -> Tuple[torch.Tensor, torch.Tensor]: # torch.compile will fuse all of the pointwise ops below # into a single kernel, making it very fast - org_vocab_mask = (input_ >= org_vocab_start_index) & (input_ < - org_vocab_end_index) + org_vocab_mask = (input_ >= org_vocab_start_index) & (input_ + < org_vocab_end_index) added_vocab_mask = (input_ >= added_vocab_start_index) & ( input_ < added_vocab_end_index) added_offset = added_vocab_start_index - ( diff --git a/fastvideo/v1/models/dits/hunyuanvideo.py b/fastvideo/v1/models/dits/hunyuanvideo.py index 9aa16a698..c3b264e9d 100644 --- a/fastvideo/v1/models/dits/hunyuanvideo.py +++ b/fastvideo/v1/models/dits/hunyuanvideo.py @@ -621,7 +621,7 @@ class HunyuanVideoTransformer3DModel(BaseDiT): # Guidance modulation self.guidance_in = (TimestepEmbedder( self.hidden_size, act_layer="silu", dtype=dtype) - if guidance_embeds else None) + if self.guidance_embeds else None) # Double blocks self.double_blocks = nn.ModuleList([ @@ -686,9 +686,9 @@ class HunyuanVideoTransformer3DModel(BaseDiT): text_states_2 = encoder_hidden_states[1] # Get spatial dimensions - _, _, ot, oh, ow = x.shape + _, _, ot, oh, ow = x.shape # codespell:ignore tt, th, tw = ( - ot // self.patch_size[0], + ot // self.patch_size[0], # codespell:ignore oh // self.patch_size[1], ow // self.patch_size[2], ) @@ -707,7 +707,7 @@ class HunyuanVideoTransformer3DModel(BaseDiT): vec = vec + self.vector_in(text_states_2) # Add guidance modulation if needed - if self.guidance_embeds and guidance is not None: + if self.guidance_in and guidance is not None: vec = vec + self.guidance_in(guidance) # Embed image and text img = self.img_in(img) diff --git a/fastvideo/v1/models/encoders/clip.py b/fastvideo/v1/models/encoders/clip.py index 186c0b75d..5fbf77c16 100644 --- a/fastvideo/v1/models/encoders/clip.py +++ b/fastvideo/v1/models/encoders/clip.py @@ -438,8 +438,9 @@ class CLIPTextTransformer(nn.Module): pooled_output = last_hidden_state[ torch.arange(last_hidden_state.shape[0], device=last_hidden_state.device), - input_ids.to(dtype=torch.int, device=last_hidden_state.device - ).argmax(dim=-1), ] + input_ids.to(dtype=torch.int, device=last_hidden_state.device). + argmax(dim=-1), + ] else: # The config gets updated `eos_token_id` from PR #24773 (so the use of exta new tokens is possible) pooled_output = last_hidden_state[ @@ -448,7 +449,8 @@ class CLIPTextTransformer(nn.Module): # We need to get the first position of `eos_token_id` value (`pad_token_ids` might equal to `eos_token_id`) # Note: we assume each sequence (along batch dim.) contains an `eos_token_id` (e.g. prepared by the tokenizer) (input_ids.to(dtype=torch.int, device=last_hidden_state.device - ) == self.eos_token_id).int().argmax(dim=-1), ] + ) == self.eos_token_id).int().argmax(dim=-1), + ] if not return_dict: return (last_hidden_state, pooled_output) + encoder_outputs[1:] diff --git a/fastvideo/v1/models/loader/fsdp_load.py b/fastvideo/v1/models/loader/fsdp_load.py index d742cd538..5d10a8738 100644 --- a/fastvideo/v1/models/loader/fsdp_load.py +++ b/fastvideo/v1/models/loader/fsdp_load.py @@ -148,7 +148,7 @@ def shard_model( reshard_after_forward (bool): Whether to reshard parameters and buffers after the forward pass. Setting this to True corresponds to the FULL_SHARD sharding strategy from FSDP1, while setting it to False corresponds to the SHARD_GRAD_OP sharding strategy. - dp_mesh (Optional[DeviceMesh]): Device mesh to use for FSDP sharding under mutliple parallelism. + dp_mesh (Optional[DeviceMesh]): Device mesh to use for FSDP sharding under multiple parallelism. Default to None. Raises: diff --git a/fastvideo/v1/models/registry.py b/fastvideo/v1/models/registry.py index df2f45eb9..e80e1e1c2 100644 --- a/fastvideo/v1/models/registry.py +++ b/fastvideo/v1/models/registry.py @@ -286,7 +286,8 @@ class _ModelRegistry: ModelRegistry = _ModelRegistry({ - model_arch: _LazyRegisteredModel( + model_arch: + _LazyRegisteredModel( module_name=f"fastvideo.v1.models.{component_name}.{mod_relname}", component_name=component_name, class_name=cls_name, diff --git a/fastvideo/v1/models/vaes/hunyuanvae.py b/fastvideo/v1/models/vaes/hunyuanvae.py index b6de1b617..0d5fd4400 100644 --- a/fastvideo/v1/models/vaes/hunyuanvae.py +++ b/fastvideo/v1/models/vaes/hunyuanvae.py @@ -566,10 +566,9 @@ class HunyuanVideoEncoder3D(nn.Module): if temporal_compression_ratio == 4: add_spatial_downsample = bool(i < num_spatial_downsample_layers) - add_time_downsample = bool( - i >= - (len(block_out_channels) - 1 - num_time_downsample_layers) - and not is_final_block) + add_time_downsample = bool(i >= (len(block_out_channels) - 1 - + num_time_downsample_layers) + and not is_final_block) elif temporal_compression_ratio == 8: add_spatial_downsample = bool(i < num_spatial_downsample_layers) add_time_downsample = bool(i < num_time_downsample_layers) diff --git a/fastvideo/v1/pipelines/README.md b/fastvideo/v1/pipelines/README.md index 6b1e0f6da..5bc734c0b 100644 --- a/fastvideo/v1/pipelines/README.md +++ b/fastvideo/v1/pipelines/README.md @@ -133,4 +133,4 @@ pipeline_cls, _ = PipelineRegistry.resolve_pipeline_cls("YourCustomPipeline") # Initialize and use the pipeline pipeline = pipeline_cls(...) result = pipeline(...) -``` \ No newline at end of file +``` diff --git a/fastvideo/v1/tests/conftest.py b/fastvideo/v1/tests/conftest.py index 6f1e88b03..eaac2668e 100644 --- a/fastvideo/v1/tests/conftest.py +++ b/fastvideo/v1/tests/conftest.py @@ -1,6 +1,10 @@ import pytest import torch.distributed as dist -from fastvideo.v1.distributed import init_distributed_environment, initialize_model_parallel, destroy_model_parallel + +from fastvideo.v1.distributed import (destroy_model_parallel, + init_distributed_environment, + initialize_model_parallel) + @pytest.fixture(scope="function") def distributed_setup(): @@ -10,15 +14,15 @@ def distributed_setup(): This ensures proper cleanup even if tests fail. """ init_distributed_environment(world_size=1, - rank=0, - distributed_init_method="env://", - local_rank=0, - backend="nccl") + rank=0, + distributed_init_method="env://", + local_rank=0, + backend="nccl") initialize_model_parallel(tensor_model_parallel_size=1, - sequence_model_parallel_size=1, - backend="nccl") + sequence_model_parallel_size=1, + backend="nccl") yield - + if dist.is_initialized(): destroy_model_parallel() dist.destroy_process_group() diff --git a/fastvideo/v1/tests/encoders/test_clip_encoder.py b/fastvideo/v1/tests/encoders/test_clip_encoder.py index 73b07579c..0c8bf1be0 100644 --- a/fastvideo/v1/tests/encoders/test_clip_encoder.py +++ b/fastvideo/v1/tests/encoders/test_clip_encoder.py @@ -1,17 +1,18 @@ # TODO: check if correct -import pytest -from fastvideo.models.hunyuan.text_encoder import load_text_encoder, load_tokenizer import os -import torch -import torch.nn as nn + import numpy as np -from fastvideo.v1.logger import init_logger +import pytest +import torch from transformers import AutoConfig -from fastvideo.v1.distributed import init_distributed_environment, initialize_model_parallel + +from fastvideo.models.hunyuan.text_encoder import (load_text_encoder, + load_tokenizer) # from fastvideo.v1.models.hunyuan.text_encoder import load_text_encoder, load_tokenizer from fastvideo.v1.forward_context import set_forward_context -from fastvideo.v1.utils import maybe_download_model from fastvideo.v1.inference_args import InferenceArgs +from fastvideo.v1.logger import init_logger +from fastvideo.v1.utils import maybe_download_model logger = init_logger(__name__) @@ -23,9 +24,12 @@ torch.manual_seed(42) np.random.seed(42) BASE_MODEL_PATH = "hunyuanvideo-community/HunyuanVideo" -MODEL_PATH = maybe_download_model(BASE_MODEL_PATH, local_dir=os.path.join("data", BASE_MODEL_PATH)) +MODEL_PATH = maybe_download_model(BASE_MODEL_PATH, + local_dir=os.path.join( + "data", BASE_MODEL_PATH)) TEXT_ENCODER_PATH = os.path.join(MODEL_PATH, "text_encoder_2") + @pytest.mark.usefixtures("distributed_setup") def test_clip_encoder(): """ @@ -37,10 +41,11 @@ def test_clip_encoder(): - Load models with the same weights and parameters - Produce nearly identical outputs for the same input prompts """ - args = InferenceArgs(model_path="openai/clip-vit-large-patch14", precision="float16") + args = InferenceArgs(model_path="openai/clip-vit-large-patch14", + precision="float16") device = torch.device("cuda:0" if torch.cuda.is_available() else "cpu") - logger.info(f"Loading models from {args.model_path}") + logger.info("Loading models from %s", args.model_path) # config = json.load(open(os.path.join(model_path, "config.json"))) @@ -72,8 +77,8 @@ def test_clip_encoder(): params2 = dict(model2.named_parameters()) # Check number of parameters - logger.info(f"Model1 has {len(params1)} parameters") - logger.info(f"Model2 has {len(params2)} parameters") + logger.info("Model1 has %d parameters", len(params1)) + logger.info("Model2 has %d parameters", len(params2)) # Compare a few key parameters @@ -103,29 +108,23 @@ def test_clip_encoder(): with torch.no_grad(): for prompt in prompts: - logger.info(f"Testing prompt: '{prompt}'") + logger.info("Testing prompt: '%s'", prompt) # Tokenize the prompt - tokens = tokenizer( - prompt, - return_tensors="pt" - ).to(device) + tokens = tokenizer(prompt, return_tensors="pt").to(device) # Get embeddings from our implementation - outputs1 = model1( - input_ids=tokens.input_ids, - output_hidden_states=True - ) - - logger.info(f"Testing model2") + outputs1 = model1(input_ids=tokens.input_ids, + output_hidden_states=True) + + logger.info("Testing model2") print("--------------------------------") # Get embeddings from HuggingFace implementation with set_forward_context(current_timestep=0, attn_metadata=None): outputs2 = model2( input_ids=tokens.input_ids, # attention_mask=tokens.attention_mask, - output_hidden_states=True - ) - + output_hidden_states=True) + # Compare last hidden states last_hidden_state1 = outputs1.last_hidden_state[ tokens.attention_mask == 1] @@ -142,12 +141,10 @@ def test_clip_encoder(): mean_diff_hidden = torch.mean( torch.abs(last_hidden_state1 - last_hidden_state2)) - logger.info( - f"Maximum difference in last hidden states: {max_diff_hidden.item()}" - ) - logger.info( - f"Mean difference in last hidden states: {mean_diff_hidden.item()}" - ) + logger.info("Maximum difference in last hidden states: %f", + max_diff_hidden.item()) + logger.info("Mean difference in last hidden states: %f", + mean_diff_hidden.item()) # Compare pooler outputs pooler_output1 = outputs1.pooler_output @@ -161,13 +158,12 @@ def test_clip_encoder(): mean_diff_pooler = torch.mean( torch.abs(pooler_output1 - pooler_output2)) - logger.info( - f"Maximum difference in pooler outputs: {max_diff_pooler.item()}" - ) - logger.info( - f"Mean difference in pooler outputs: {mean_diff_pooler.item()}") - - # Check if outputs are similar (allowing for small numerical differences) + logger.info("Maximum difference in pooler outputs: %f", + max_diff_pooler.item()) + logger.info("Mean difference in pooler outputs: %f", + mean_diff_pooler.item()) + + # Check if outputs are similar (allowing for small numerical differences) assert mean_diff_hidden < 1e-2, \ f"Hidden states differ significantly: mean diff = {mean_diff_hidden.item()}" assert mean_diff_pooler < 1e-2, \ diff --git a/fastvideo/v1/tests/encoders/test_llama_encoder.py b/fastvideo/v1/tests/encoders/test_llama_encoder.py index 62dfc819e..f812ab0b7 100644 --- a/fastvideo/v1/tests/encoders/test_llama_encoder.py +++ b/fastvideo/v1/tests/encoders/test_llama_encoder.py @@ -1,15 +1,17 @@ -import pytest -from fastvideo.models.hunyuan.text_encoder import load_text_encoder, load_tokenizer import os -import torch -import torch.nn as nn + import numpy as np -from fastvideo.v1.logger import init_logger +import pytest +import torch from transformers import AutoConfig -from fastvideo.v1.models.loader.component_loader import TextEncoderLoader + +from fastvideo.models.hunyuan.text_encoder import (load_text_encoder, + load_tokenizer) from fastvideo.v1.forward_context import set_forward_context -from fastvideo.v1.utils import maybe_download_model from fastvideo.v1.inference_args import InferenceArgs +from fastvideo.v1.logger import init_logger +from fastvideo.v1.models.loader.component_loader import TextEncoderLoader +from fastvideo.v1.utils import maybe_download_model logger = init_logger(__name__) @@ -21,10 +23,13 @@ torch.manual_seed(42) np.random.seed(42) BASE_MODEL_PATH = "hunyuanvideo-community/HunyuanVideo" -MODEL_PATH = maybe_download_model(BASE_MODEL_PATH, local_dir=os.path.join('data', BASE_MODEL_PATH)) +MODEL_PATH = maybe_download_model(BASE_MODEL_PATH, + local_dir=os.path.join( + 'data', BASE_MODEL_PATH)) TEXT_ENCODER_PATH = os.path.join(MODEL_PATH, "text_encoder") TOKENIZER_PATH = os.path.join(MODEL_PATH, "tokenizer") + @pytest.mark.usefixtures("distributed_setup") def test_llama_encoder(): """ @@ -36,23 +41,22 @@ def test_llama_encoder(): - Load models with the same weights and parameters - Produce nearly identical outputs for the same input prompts """ - args = InferenceArgs(model_path="meta-llama/Llama-2-7b-hf", precision="float16") + args = InferenceArgs(model_path="meta-llama/Llama-2-7b-hf", + precision="float16") device = torch.device("cuda:0" if torch.cuda.is_available() else "cpu") # Initialize the two model implementations - logger.info(f"Loading models from {args.model_path}") + logger.info("Loading models from %s", args.model_path) hf_config = AutoConfig.from_pretrained(TEXT_ENCODER_PATH) print(hf_config) # Load our implementation using the loader from text_encoder/__init__.py - model1, _ = load_text_encoder( - text_encoder_type="llm", - text_encoder_precision='fp16', - text_encoder_path=TEXT_ENCODER_PATH, - logger=logger, - device=device - ) + model1, _ = load_text_encoder(text_encoder_type="llm", + text_encoder_precision='fp16', + text_encoder_path=TEXT_ENCODER_PATH, + logger=logger, + device=device) loader = TextEncoderLoader() args.device_str = "cuda:0" device = torch.device(args.device_str) @@ -69,8 +73,8 @@ def test_llama_encoder(): params2 = dict(model2.named_parameters()) # Check number of parameters - logger.info(f"Model1 has {len(params1)} parameters") - logger.info(f"Model2 has {len(params2)} parameters") + logger.info("Model1 has %d parameters", len(params1)) + logger.info("Model2 has %d parameters", len(params2)) # Compare a few key parameters weight_diffs = [] @@ -85,10 +89,10 @@ def test_llama_encoder(): # for (name1, param1), (name2, param2) in zip( # sorted(params1.items()), sorted(params2.items()) # ): - for l in range(hf_config.num_hidden_layers): + for layer_idx in range(hf_config.num_hidden_layers): for w in weights: - name1 = w.format(l) - name2 = w.format(l) + name1 = w.format(layer_idx) + name2 = w.format(layer_idx) p1 = params1[name1] p2 = params2[name2] # print(type(p2)) @@ -102,7 +106,7 @@ def test_llama_encoder(): weight_diffs.append((name1, name2, max_diff, mean_diff)) # logger.info(f" Max diff: {max_diff}, Mean diff: {mean_diff}") except Exception as e: - logger.info(f"Error comparing {name1} and {name2}: {e}") + logger.info("Error comparing %s and %s: %s", name1, name2, e) tokenizer, _ = load_tokenizer(tokenizer_type="llm", tokenizer_path=TOKENIZER_PATH, @@ -119,33 +123,26 @@ def test_llama_encoder(): with torch.no_grad(): for prompt in prompts: - logger.info(f"Testing prompt: '{prompt}'") + logger.info("Testing prompt: '%s'", prompt) # Tokenize the prompt - tokens = tokenizer( - prompt, - return_tensors="pt" - ).to(device) - + tokens = tokenizer(prompt, return_tensors="pt").to(device) + # Get outputs from our implementation # filter out padding input_ids # tokens.input_ids = tokens.input_ids[tokens.attention_mask==1] # tokens.attention_mask = tokens.attention_mask[tokens.attention_mask==1] - outputs1 = model1( - input_ids=tokens.input_ids, - output_hidden_states=True - ) + outputs1 = model1(input_ids=tokens.input_ids, + output_hidden_states=True) print("--------------------------------") - logger.info(f"Testing model2") + logger.info("Testing model2") # Get outputs from HuggingFace implementation with set_forward_context(current_timestep=0, attn_metadata=None): - outputs2 = model2( - input_ids=tokens.input_ids, - attention_mask=tokens.attention_mask, - output_hidden_states=True - ) - + outputs2 = model2(input_ids=tokens.input_ids, + attention_mask=tokens.attention_mask, + output_hidden_states=True) + # Compare last hidden states last_hidden_state1 = outputs1.last_hidden_state[ tokens.attention_mask == 1] @@ -160,14 +157,12 @@ def test_llama_encoder(): mean_diff_hidden = torch.mean( torch.abs(last_hidden_state1 - last_hidden_state2)) - logger.info( - f"Maximum difference in last hidden states: {max_diff_hidden.item()}" - ) - logger.info( - f"Mean difference in last hidden states: {mean_diff_hidden.item()}" - ) + logger.info("Maximum difference in last hidden states: %f", + max_diff_hidden.item()) + logger.info("Mean difference in last hidden states: %f", + mean_diff_hidden.item()) - # Check if outputs are similar (allowing for small numerical differences) + # Check if outputs are similar (allowing for small numerical differences) assert mean_diff_hidden < 1e-2, \ f"Hidden states differ significantly: mean diff = {mean_diff_hidden.item()}" assert max_diff_hidden < 1e-1, \ diff --git a/fastvideo/v1/tests/old_tests/test_hunyuan_vae.py b/fastvideo/v1/tests/old_tests/test_hunyuan_vae.py index f963ed724..99ac4151a 100644 --- a/fastvideo/v1/tests/old_tests/test_hunyuan_vae.py +++ b/fastvideo/v1/tests/old_tests/test_hunyuan_vae.py @@ -1,15 +1,17 @@ -from fastvideo.v1.models.vaes.hunyuanvae import AutoencoderKLHunyuanVideo as MyHunyuanVAE -from diffusers import AutoencoderKLHunyuanVideo as DiffusersHunyuanVAE - +import argparse +import json import os + +import numpy as np import torch import torch.nn as nn -import argparse -import numpy as np -import json -from fastvideo.v1.logger import init_logger +from diffusers import AutoencoderKLHunyuanVideo as DiffusersHunyuanVAE from safetensors.torch import load_file +from fastvideo.v1.logger import init_logger +from fastvideo.v1.models.vaes.hunyuanvae import ( + AutoencoderKLHunyuanVideo as MyHunyuanVAE) + logger = init_logger(__name__) diff --git a/fastvideo/v1/tests/old_tests/test_hunyuanvideo.py b/fastvideo/v1/tests/old_tests/test_hunyuanvideo.py index c530f5e76..b8a882ed5 100644 --- a/fastvideo/v1/tests/old_tests/test_hunyuanvideo.py +++ b/fastvideo/v1/tests/old_tests/test_hunyuanvideo.py @@ -1,24 +1,25 @@ -from itertools import chain +import argparse import os +from itertools import chain + +import numpy as np import torch import torch.nn as nn -import argparse -import numpy as np - -from fastvideo.v1.logger import init_logger -from fastvideo.v1.distributed.parallel_state import ( - init_distributed_environment, initialize_model_parallel, - get_sequence_model_parallel_rank, get_sequence_model_parallel_world_size, - destroy_model_parallel, destroy_distributed_environment, - cleanup_dist_env_and_memory) - -from fastvideo.v1.utils.parallel_states import initialize_sequence_parallel_state -from torch.distributed._composable.fsdp import CPUOffloadPolicy, fully_shard from torch.distributed.device_mesh import init_device_mesh + +from fastvideo.v1.distributed.parallel_state import ( + cleanup_dist_env_and_memory, destroy_distributed_environment, + destroy_model_parallel, get_sequence_model_parallel_rank, + get_sequence_model_parallel_world_size, init_distributed_environment, + initialize_model_parallel) +from fastvideo.v1.logger import init_logger +from fastvideo.v1.models.dits.hunyuanvideo import ( + HunyuanVideoTransformer3DModel as HunyuanVideoDit) +from fastvideo.v1.models.hunyuan.modules.models import ( + HYVideoDiffusionTransformer) from fastvideo.v1.models.loader.fsdp_load import shard_model -from fastvideo.v1.models.dits.hunyuanvideo import HunyuanVideoTransformer3DModel as HunyuanVideoDit -from fastvideo.v1.models.hunyuan.modules.models import HYVideoDiffusionTransformer -from fastvideo.v1.models.hunyuan_hf.modeling_hunyuan import HunyuanVideoTransformer3DModel +from fastvideo.v1.utils.parallel_states import ( + initialize_sequence_parallel_state) logger = init_logger(__name__) diff --git a/fastvideo/v1/tests/old_tests/test_hunyuanvideo_load.py b/fastvideo/v1/tests/old_tests/test_hunyuanvideo_load.py index ab87ccbf2..fa471e777 100644 --- a/fastvideo/v1/tests/old_tests/test_hunyuanvideo_load.py +++ b/fastvideo/v1/tests/old_tests/test_hunyuanvideo_load.py @@ -1,20 +1,22 @@ -import os -import torch import argparse - -from fastvideo.v1.logger import init_logger -from fastvideo.v1.distributed.parallel_state import ( - init_distributed_environment, initialize_model_parallel, - get_sequence_model_parallel_rank, get_sequence_model_parallel_world_size, - destroy_model_parallel, destroy_distributed_environment, - cleanup_dist_env_and_memory) -import json -from fastvideo.v1.models.dits.hunyuanvideo import HunyuanVideoTransformer3DModel as HunyuanVideoDit -from fastvideo.models.hunyuan.modules.models import HUNYUAN_VIDEO_CONFIG -from fastvideo.models.hunyuan.modules.models import HYVideoDiffusionTransformer -from fastvideo.v1.models.loader.fsdp_load import load_fsdp_model -from fastvideo.utils.parallel_states import initialize_sequence_parallel_state import glob +import json +import os + +import torch + +from fastvideo.models.hunyuan.modules.models import ( + HUNYUAN_VIDEO_CONFIG, HYVideoDiffusionTransformer) +from fastvideo.utils.parallel_states import initialize_sequence_parallel_state +from fastvideo.v1.distributed.parallel_state import ( + cleanup_dist_env_and_memory, destroy_distributed_environment, + destroy_model_parallel, get_sequence_model_parallel_rank, + get_sequence_model_parallel_world_size, init_distributed_environment, + initialize_model_parallel) +from fastvideo.v1.logger import init_logger +from fastvideo.v1.models.dits.hunyuanvideo import ( + HunyuanVideoTransformer3DModel as HunyuanVideoDit) +from fastvideo.v1.models.loader.fsdp_load import load_fsdp_model logger = init_logger(__name__) @@ -48,8 +50,7 @@ def test_hunyuanvideo_distributed(): torch.cuda.set_device(f"cuda:{local_rank}") # Initialize tensor model parallel groups initialize_model_parallel( - sequence_model_parallel_size=args.sequence_model_parallel_size - ) + sequence_model_parallel_size=args.sequence_model_parallel_size) initialize_sequence_parallel_state(args.sequence_model_parallel_size) # Get tensor parallel info sp_rank = get_sequence_model_parallel_rank() @@ -61,8 +62,7 @@ def test_hunyuanvideo_distributed(): # load data/hunyuanvideo_community/transformer/config.json with open( - "data/hunyuanvideo-community/HunyuanVideo/transformer/config.json", - "r") as f: + "data/hunyuanvideo-community/HunyuanVideo/transformer/config.json") as f: config = json.load(f) # remove "_class_name": "HunyuanVideoTransformer3DModel", "_diffusers_version": "0.32.0.dev0", # TODO: write normalize config function @@ -154,14 +154,14 @@ def test_hunyuanvideo_distributed(): with torch.no_grad(): # Run inference on model1 with torch.amp.autocast(device_type="cuda", dtype=torch.bfloat16): - logger.info(f"Running inference on model1") + logger.info("Running inference on model1") output1 = model1( hidden_states=hidden_states, encoder_hidden_states=encoder_hidden_states, timestep=timestep, ) logger.info("Model 1 inference completed") - + # Run inference on model2 output2, _ = model2( hidden_states=hidden_states, @@ -184,13 +184,12 @@ def test_hunyuanvideo_distributed(): logger.info(f"Model 2 weight mean: {weight_mean_model2}") weight_mean_diff = abs(weight_mean_model1 - weight_mean_model2) logger.info(f"Weight mean difference: {weight_mean_diff}") - # mean diff mean_diff = torch.mean(torch.abs(output1 - output2)) assert mean_diff < 1e-2, f"Mean difference between outputs: {mean_diff.item()}" - - # diff sum + + # diff sum diff_sum = torch.sum(torch.abs(output1 - output2)) logger.info(f"Diff sum between outputs: {diff_sum.item()}") diff --git a/fastvideo/v1/tests/old_tests/test_t5_encoder.py b/fastvideo/v1/tests/old_tests/test_t5_encoder.py index 69656ee35..3b54954df 100644 --- a/fastvideo/v1/tests/old_tests/test_t5_encoder.py +++ b/fastvideo/v1/tests/old_tests/test_t5_encoder.py @@ -1,57 +1,70 @@ -from transformers import UMT5EncoderModel, AutoTokenizer, AutoConfig -import os -import torch -import torch.nn as nn import argparse + import numpy as np +import torch +from transformers import AutoConfig, AutoTokenizer, UMT5EncoderModel + +from fastvideo.v1.distributed import (init_distributed_environment, + initialize_model_parallel) from fastvideo.v1.logger import init_logger -from fastvideo.v1.distributed import init_distributed_environment, initialize_model_parallel logger = init_logger(__name__) + def setup_args(): parser = argparse.ArgumentParser(description='T5 Encoder Test') parser.add_argument('--model_path', type=str, default="google/umt5-xxl") - parser.add_argument('--precision', type=str, default="float32", - help='Precision to use for the model (float32, float16, bfloat16)') + parser.add_argument( + '--precision', + type=str, + default="float32", + help='Precision to use for the model (float32, float16, bfloat16)') return parser.parse_args() + def test_t5_encoder(): - init_distributed_environment(world_size=1, rank=0, distributed_init_method="env://", local_rank=0, backend="nccl") - initialize_model_parallel(tensor_model_parallel_size=1, sequence_model_parallel_size=1, backend="nccl") + init_distributed_environment(world_size=1, + rank=0, + distributed_init_method="env://", + local_rank=0, + backend="nccl") + initialize_model_parallel(tensor_model_parallel_size=1, + sequence_model_parallel_size=1, + backend="nccl") args = setup_args() - + # Set fixed random seed for reproducibility torch.manual_seed(42) np.random.seed(42) - + device = torch.device("cuda:0" if torch.cuda.is_available() else "cpu") - + # Initialize the two model implementations model_path = "/workspace/data/Wan2.1-T2V-1.3B-Diffusers/text_encoder" tokenizer_path = "/workspace/data/Wan2.1-T2V-1.3B-Diffusers/tokenizer" hf_config = AutoConfig.from_pretrained(model_path) print(hf_config) - precision = torch.float16 # It must be float16 because the weight loader is in float16 + precision = torch.float16 # It must be float16 because the weight loader is in float16 # Load our implementation using the loader from text_encoder/__init__.py - model1 = UMT5EncoderModel.from_pretrained(model_path).to(precision).to(device).eval() + model1 = UMT5EncoderModel.from_pretrained(model_path).to(precision).to( + device).eval() tokenizer = AutoTokenizer.from_pretrained(tokenizer_path) from fastvideo.v1.models.loader.component_loader import TextEncoderLoader loader = TextEncoderLoader() model2 = loader.load_model(model_path, hf_config, device) - + # Convert to float16 and move to device model2 = model2.to(precision) model2 = model2.to(device) model2.eval() - + # Sanity check weights between the two models logger.info("Comparing model weights for sanity check...") params1 = dict(model1.named_parameters()) params2 = dict(model2.named_parameters()) - + # Check number of parameters logger.info(f"Model1 has {len(params1)} parameters") logger.info(f"Model2 has {len(params2)} parameters") @@ -80,76 +93,80 @@ def test_t5_encoder(): logger.info(f" Max diff: {max_diff}, Mean diff: {mean_diff}") except Exception as e: logger.info(f"Error comparing {name1} and {name2}: {e}") - + total_params = sum(p.numel() for p in model1.parameters()) - weight_sum_model1 = sum(p.to(torch.float64).sum().item() for p in model1.parameters()) + weight_sum_model1 = sum( + p.to(torch.float64).sum().item() for p in model1.parameters()) weight_mean_model1 = weight_sum_model1 / total_params print("Model 1 Weight Sum: ", weight_sum_model1) print("Model 1 Weight Mean: ", weight_mean_model1) total_params = sum(p.numel() for p in model2.parameters()) - weight_sum_model2 = sum(p.to(torch.float64).sum().item() for p in model2.parameters()) + weight_sum_model2 = sum( + p.to(torch.float64).sum().item() for p in model2.parameters()) # Also calculate mean for more stable comparison weight_mean_model2 = weight_sum_model2 / total_params print("Model 2 Weight Sum: ", weight_sum_model2) print("Model 2 Weight Mean: ", weight_mean_model2) - + # Test with some sample prompts prompts = [ - "Once upon a time", - "The quick brown fox jumps over", + "Once upon a time", "The quick brown fox jumps over", "In a galaxy far, far away" ] - + logger.info("Testing T5 encoder with sample prompts") - + with torch.no_grad(): for prompt in prompts: logger.info(f"Testing prompt: '{prompt}'") - + # Tokenize the prompt - tokens = tokenizer( - prompt, - padding="max_length", - max_length=512, - truncation=True, - return_tensors="pt" - ).to(device) - + tokens = tokenizer(prompt, + padding="max_length", + max_length=512, + truncation=True, + return_tensors="pt").to(device) + # Get outputs from our implementation - # filter out padding input_ids + # filter out padding input_ids # tokens.input_ids = tokens.input_ids[tokens.attention_mask==1] # tokens.attention_mask = tokens.attention_mask[tokens.attention_mask==1] - outputs1 = model1( - input_ids=tokens.input_ids, - attention_mask=tokens.attention_mask, - output_hidden_states=True - ).last_hidden_state + outputs1 = model1(input_ids=tokens.input_ids, + attention_mask=tokens.attention_mask, + output_hidden_states=True).last_hidden_state print("--------------------------------") - logger.info(f"Testing model2") - + logger.info("Testing model2") + # Get outputs from HuggingFace implementation outputs2 = model2( input_ids=tokens.input_ids, attention_mask=tokens.attention_mask, ) - + # Compare last hidden states - last_hidden_state1 = outputs1[tokens.attention_mask==1] - last_hidden_state2 = outputs2[tokens.attention_mask==1] - + last_hidden_state1 = outputs1[tokens.attention_mask == 1] + last_hidden_state2 = outputs2[tokens.attention_mask == 1] + assert last_hidden_state1.shape == last_hidden_state2.shape, \ f"Hidden state shapes don't match: {last_hidden_state1.shape} vs {last_hidden_state2.shape}" - - max_diff_hidden = torch.max(torch.abs(last_hidden_state1 - last_hidden_state2)) - mean_diff_hidden = torch.mean(torch.abs(last_hidden_state1 - last_hidden_state2)) - - logger.info(f"Maximum difference in last hidden states: {max_diff_hidden.item()}") - logger.info(f"Mean difference in last hidden states: {mean_diff_hidden.item()}") - - - logger.info("Test passed! Both T5 encoder implementations produce similar outputs.") + + max_diff_hidden = torch.max( + torch.abs(last_hidden_state1 - last_hidden_state2)) + mean_diff_hidden = torch.mean( + torch.abs(last_hidden_state1 - last_hidden_state2)) + + logger.info( + f"Maximum difference in last hidden states: {max_diff_hidden.item()}" + ) + logger.info( + f"Mean difference in last hidden states: {mean_diff_hidden.item()}" + ) + + logger.info( + "Test passed! Both T5 encoder implementations produce similar outputs.") logger.info("Test completed successfully") + if __name__ == "__main__": test_t5_encoder() diff --git a/fastvideo/v1/tests/old_tests/test_wan_vae.py b/fastvideo/v1/tests/old_tests/test_wan_vae.py index 7e0161747..415f2f464 100644 --- a/fastvideo/v1/tests/old_tests/test_wan_vae.py +++ b/fastvideo/v1/tests/old_tests/test_wan_vae.py @@ -1,22 +1,22 @@ -from fastvideo.v1.models.vaes.wanvae import AutoencoderKLWan as MyWanVAE -from diffusers import AutoencoderKLWan - -import os -import gc -import torch -import torch.nn as nn -import argparse -import numpy as np import json -from fastvideo.v1.logger import init_logger +import os + +import numpy as np +import torch +from diffusers import AutoencoderKLWan from safetensors.torch import load_file + +from fastvideo.v1.logger import init_logger +from fastvideo.v1.models.vaes.wanvae import AutoencoderKLWan as MyWanVAE + logger = init_logger(__name__) -def test_wan_vae(): + +def test_wan_vae(): # Set fixed random seed for reproducibility torch.manual_seed(42) np.random.seed(42) - + device = torch.device("cuda:0") # Initialize the two model implementations path = "/workspace/data/Wan2.1-T2V-1.3B-Diffusers/vae" @@ -24,27 +24,22 @@ def test_wan_vae(): config = json.load(open(config_path)) config.pop("_class_name") config.pop("_diffusers_version") - model1 = MyWanVAE( - **config - ).to(torch.bfloat16) - - model2 = AutoencoderKLWan( - **config - ).to(torch.bfloat16) - - loaded = load_file(os.path.join(path, "diffusion_pytorch_model.safetensors")) + model1 = MyWanVAE(**config).to(torch.bfloat16) + + model2 = AutoencoderKLWan(**config).to(torch.bfloat16) + + loaded = load_file(os.path.join(path, + "diffusion_pytorch_model.safetensors")) model1.load_state_dict(loaded) model2.load_state_dict(loaded) - # Set both models to eval mode model1.eval() model2.eval() - + # Move to GPU model1 = model1.to(device) model2 = model2.to(device) - # model1.enable_tiling( # tile_sample_min_height=32, @@ -54,14 +49,26 @@ def test_wan_vae(): # tile_sample_stride_width=16, # tile_sample_stride_num_frames=4 # ) - + # Create identical inputs for both models batch_size = 1 - + # Video input [B, C, T, H, W] - input_tensor = torch.randn(batch_size, 3, 81, 32, 32, device=device, dtype=torch.bfloat16) - latent_tensor = torch.randn(batch_size, 16, 21, 32, 32, device=device, dtype=torch.bfloat16) - + input_tensor = torch.randn(batch_size, + 3, + 81, + 32, + 32, + device=device, + dtype=torch.bfloat16) + latent_tensor = torch.randn(batch_size, + 16, + 21, + 32, + 32, + device=device, + dtype=torch.bfloat16) + # Disable gradients for inference with torch.no_grad(): # Test encoding @@ -75,36 +82,41 @@ def test_wan_vae(): # Check if latents are similar max_diff_encode = torch.max(torch.abs(latent1 - latent2)) mean_diff_encode = torch.mean(torch.abs(latent1 - latent2)) - logger.info(f"Maximum difference between encoded latents: {max_diff_encode.item()}") - logger.info(f"Mean difference between encoded latents: {mean_diff_encode.item()}") + logger.info( + f"Maximum difference between encoded latents: {max_diff_encode.item()}" + ) + logger.info( + f"Mean difference between encoded latents: {mean_diff_encode.item()}" + ) assert mean_diff_encode < 5e-1, f"Encoded latents differ significantly: mean diff = {mean_diff_encode.item()}" # Test decoding logger.info("Testing decoding...") latent1 = latent2 = latent_tensor - latents_mean = ( - torch.tensor(model2.config.latents_mean) - .view(1, model2.config.z_dim, 1, 1, 1) - .to(latent2.device, latent2.dtype) - ) - latents_std = 1.0 / torch.tensor(model2.config.latents_std).view(1, model2.config.z_dim, 1, 1, 1).to( - latent2.device, latent2.dtype - ) + latents_mean = (torch.tensor(model2.config.latents_mean).view( + 1, model2.config.z_dim, 1, 1, 1).to(latent2.device, latent2.dtype)) + latents_std = 1.0 / torch.tensor(model2.config.latents_std).view( + 1, model2.config.z_dim, 1, 1, 1).to(latent2.device, latent2.dtype) latent2 = latent2 / latents_std + latents_mean output1 = model1.decode(latent1) output2 = model2.decode(latent2).sample # Check if outputs have the same shape assert output1.shape == output2.shape, f"Output shapes don't match: {output1.shape} vs {output2.shape}" - + # Check if outputs are similar max_diff_decode = torch.max(torch.abs(output1 - output2)) mean_diff_decode = torch.mean(torch.abs(output1 - output2)) - logger.info(f"Maximum difference between decoded outputs: {max_diff_decode.item()}") - logger.info(f"Mean difference between decoded outputs: {mean_diff_decode.item()}") + logger.info( + f"Maximum difference between decoded outputs: {max_diff_decode.item()}" + ) + logger.info( + f"Mean difference between decoded outputs: {mean_diff_decode.item()}" + ) assert mean_diff_decode < 1e-1, f"Decoded outputs differ significantly: mean diff = {mean_diff_decode.item()}" - - logger.info("Test passed! Both VAE implementations produce similar outputs.") + logger.info( + "Test passed! Both VAE implementations produce similar outputs.") logger.info("Test completed successfully") + if __name__ == "__main__": - test_wan_vae() \ No newline at end of file + test_wan_vae() diff --git a/fastvideo/v1/tests/old_tests/tp_example.py b/fastvideo/v1/tests/old_tests/tp_example.py index d555f75cd..de0643c78 100644 --- a/fastvideo/v1/tests/old_tests/tp_example.py +++ b/fastvideo/v1/tests/old_tests/tp_example.py @@ -1,18 +1,16 @@ +import argparse import os + import torch import torch.nn as nn -import torch.distributed as dist -import argparse -from fastvideo.v1.logger import init_logger from fastvideo.v1.distributed.parallel_state import ( - init_distributed_environment, initialize_model_parallel, - get_tensor_model_parallel_rank, get_tensor_model_parallel_world_size, - destroy_model_parallel, destroy_distributed_environment, - cleanup_dist_env_and_memory) -from fastvideo.v1.layers.linear import (ColumnParallelLinear, RowParallelLinear) -from fastvideo.v1.distributed.communication_op import ( - tensor_model_parallel_all_reduce, tensor_model_parallel_all_gather) + cleanup_dist_env_and_memory, destroy_distributed_environment, + destroy_model_parallel, get_tensor_model_parallel_rank, + get_tensor_model_parallel_world_size, init_distributed_environment, + initialize_model_parallel) +from fastvideo.v1.layers.linear import ColumnParallelLinear, RowParallelLinear +from fastvideo.v1.logger import init_logger logger = init_logger(__name__) diff --git a/fastvideo/v1/tests/ssim/README.md b/fastvideo/v1/tests/ssim/README.md index 6f02e6dee..f62ebb836 100644 --- a/fastvideo/v1/tests/ssim/README.md +++ b/fastvideo/v1/tests/ssim/README.md @@ -26,7 +26,7 @@ The reference videos in the `reference_videos` directory are used as part of an "fps": 24 } -#### Prompts +### Prompts 1. Will Smith casually eats noodles, his relaxed demeanor contrasting with the energetic background of a bustling street food market. The scene captures a mix of humor and authenticity. Mid-shot framing, vibrant lighting." diff --git a/fastvideo/v1/tests/ssim/compute_ssim.py b/fastvideo/v1/tests/ssim/compute_ssim.py index 357322a00..8a6647903 100644 --- a/fastvideo/v1/tests/ssim/compute_ssim.py +++ b/fastvideo/v1/tests/ssim/compute_ssim.py @@ -1,8 +1,9 @@ -import torch -import numpy as np -import os import argparse -from pytorch_msssim import ssim, ms_ssim +import os + +import numpy as np +import torch +from pytorch_msssim import ms_ssim, ssim from torchvision.io import read_video @@ -65,7 +66,7 @@ def compare_folders(reference_folder, generated_folder, use_ms_ssim=True): Compare videos with the same filename between reference_folder and generated_folder """ reference_videos = [ - f for f in os.listdir(reference_folder) if f.endswith(('.mp4')) + f for f in os.listdir(reference_folder) if f.endswith('.mp4') ] results = {} diff --git a/fastvideo/v1/tests/ssim/test_inference_similarity.py b/fastvideo/v1/tests/ssim/test_inference_similarity.py index 1854c7c23..39c29a794 100644 --- a/fastvideo/v1/tests/ssim/test_inference_similarity.py +++ b/fastvideo/v1/tests/ssim/test_inference_similarity.py @@ -1,9 +1,11 @@ -import os -import pytest import json -from fastvideo.v1.tests.ssim.compute_ssim import compute_video_ssim_torchvision -from fastvideo.v1.logger import init_logger +import os + +import pytest + from fastvideo.v1.entrypoints.cli.utils import launch_distributed +from fastvideo.v1.logger import init_logger +from fastvideo.v1.tests.ssim.compute_ssim import compute_video_ssim_torchvision logger = init_logger(__name__) @@ -131,7 +133,8 @@ def test_inference_similarity(num_inference_steps, prompt, ATTENTION_BACKEND): if not os.path.exists(reference_folder): logger.error("Reference folder missing") - raise FileNotFoundError(f"Reference video folder does not exist: {reference_folder}") + raise FileNotFoundError( + f"Reference video folder does not exist: {reference_folder}") # Find the matching reference video based on the prompt reference_video_name = None diff --git a/fastvideo/v1/utils.py b/fastvideo/v1/utils.py index 0c4dea81d..4d6d2b84d 100644 --- a/fastvideo/v1/utils.py +++ b/fastvideo/v1/utils.py @@ -11,7 +11,7 @@ import os import sys import tempfile from functools import wraps -from typing import Any, Dict, List, Type, TypeVar, Union, cast +from typing import Any, Dict, List, Optional, Type, TypeVar, Union, cast import filelock import torch @@ -380,7 +380,8 @@ def import_pynvml(): return pynvml -def maybe_download_model(model_path: str, local_dir: str = None) -> str: +def maybe_download_model(model_path: str, + local_dir: Optional[str] = None) -> str: """ Check if the model path is a Hugging Face Hub model ID and download it if needed. @@ -404,8 +405,7 @@ def maybe_download_model(model_path: str, local_dir: str = None) -> str: local_path = snapshot_download( repo_id=model_path, ignore_patterns=["*.onnx", "*.msgpack"], - local_dir=local_dir - ) + local_dir=local_dir) logger.info("Downloaded model to %s", local_path) return str(local_path) except Exception as e: diff --git a/format.sh b/format.sh deleted file mode 100644 index 73e2aaeb5..000000000 --- a/format.sh +++ /dev/null @@ -1,275 +0,0 @@ -#!/usr/bin/env bash -# YAPF formatter, adapted from ray and skypilot. -# -# Usage: -# # Do work and commit your work. - -# # Format files that differ from origin/main. -# bash format.sh - -# # Commit changed files with message 'Run yapf and ruff' -# -# -# YAPF + Clang formatter (if installed). This script formats all changed files from the last mergebase. -# You are encouraged to run this locally before pushing changes for review. - -# Cause the script to exit if a single command fails -set -eo pipefail - -# this stops git rev-parse from failing if we run this from the .git directory -builtin cd "$(dirname "${BASH_SOURCE:-$0}")" -ROOT="$(git rev-parse --show-toplevel)" -builtin cd "$ROOT" || exit 1 - -check_command() { - if ! command -v "$1" &> /dev/null; then - echo "❓❓$1 is not installed, please run \`pip install -r requirements-lint.txt\`" - exit 1 - fi -} - -check_command yapf -check_command ruff -check_command mypy -check_command codespell -check_command isort - -YAPF_VERSION=$(yapf --version | awk '{print $2}') -RUFF_VERSION=$(ruff --version | awk '{print $2}') -MYPY_VERSION=$(mypy --version | awk '{print $2}') -CODESPELL_VERSION=$(codespell --version) -ISORT_VERSION=$(isort --vn) -SPHINX_LINT_VERSION=$(sphinx-lint --version | awk '{print $2}') - -# # params: tool name, tool version, required version -tool_version_check() { - expected=$(grep "$1" requirements-lint.txt | cut -d'=' -f3) - if [[ "$2" != "$expected" ]]; then - echo "❓❓Wrong $1 version installed: $expected is required, not $2." - exit 1 - fi -} - -tool_version_check "yapf" "$YAPF_VERSION" -tool_version_check "ruff" "$RUFF_VERSION" -tool_version_check "mypy" "$MYPY_VERSION" -tool_version_check "isort" "$ISORT_VERSION" -tool_version_check "codespell" "$CODESPELL_VERSION" - -YAPF_FLAGS=( - '--recursive' - '--parallel' -) - -YAPF_EXCLUDES=( - '--exclude' 'build/**' -) - -# Format specified files -format() { - yapf --in-place "${YAPF_FLAGS[@]}" "$@" -} - -# Format files that differ from main branch. Ignores dirs that are not slated -# for autoformat yet. -format_changed() { - # The `if` guard ensures that the list of filenames is not empty, which - # could cause yapf to receive 0 positional arguments, making it hang - # waiting for STDIN. - # - # `diff-filter=ACM` and $MERGEBASE is to ensure we only format files that - # exist on both branches. - MERGEBASE="$(git merge-base origin/main HEAD)" - - if ! git diff --diff-filter=ACM --quiet --exit-code "$MERGEBASE" -- '*.py' '*.pyi' &>/dev/null; then - git diff --name-only --diff-filter=ACM "$MERGEBASE" -- '*.py' '*.pyi' | xargs -P 5 \ - yapf --in-place "${YAPF_EXCLUDES[@]}" "${YAPF_FLAGS[@]}" - fi - -} - -# Format all files -format_all() { - yapf --in-place "${YAPF_FLAGS[@]}" "${YAPF_EXCLUDES[@]}" . -} - -## This flag formats individual files. --files *must* be the first command line -## arg to use this option. -if [[ "$1" == '--files' ]]; then - format "${@:2}" - # If `--all` is passed, then any further arguments are ignored and the - # entire python directory is formatted. -elif [[ "$1" == '--all' ]]; then - format_all -else - # Format only the files that changed in last commit. - format_changed -fi -echo 'FastVideo yapf: Done' - -# Run mypy -echo 'FastVideo mypy:' - -CI=${1:-0} -PYTHON_VERSION=3.10 - -# Run mypy on specified files -run_mypy() { - # if [ "$CI" -eq 1 ]; then - # mypy --python-version "${PYTHON_VERSION}" "$@" - # else - mypy --follow-imports skip --python-version "${PYTHON_VERSION}" "$@" - # fi -} - -# Run mypy on all files -run_mypy_all() { - run_mypy fastvideo tests -} - -# Run mypy on changed files -run_mypy_changed() { - MERGEBASE="$(git merge-base origin/main HEAD)" - - if ! git diff --diff-filter=ACM --quiet --exit-code "$MERGEBASE" -- '*.py' '*.pyi' &>/dev/null; then - git diff --name-only --diff-filter=ACM "$MERGEBASE" -- '*.py' '*.pyi' | xargs \ - run_mypy - fi -} - -# Run mypy based on command line arguments -if [[ "$1" == '--files' ]]; then - run_mypy "${@:2}" -elif [[ "$1" == '--all' ]]; then - run_mypy_all -else - run_mypy_changed -fi - -echo 'FastVideo mypy: Done' - - -# If git diff returns a file that is in the skip list, the file may be checked anyway: -# https://github.com/codespell-project/codespell/issues/1915 -# Avoiding the "./" prefix and using "/**" globs for directories appears to solve the problem -CODESPELL_EXCLUDES=( - '--skip' 'tests/prompts/**,./benchmarks/sonnet.txt,*tests/lora/data/**,build/**' -) - -# check spelling of specified files -spell_check() { - codespell "$@" -} - -spell_check_all(){ - codespell --toml pyproject.toml "${CODESPELL_EXCLUDES[@]}" -} - -# Spelling check of files that differ from main branch. -spell_check_changed() { - # The `if` guard ensures that the list of filenames is not empty, which - # could cause ruff to receive 0 positional arguments, making it hang - # waiting for STDIN. - # - # `diff-filter=ACM` and $MERGEBASE is to ensure we only lint files that - # exist on both branches. - MERGEBASE="$(git merge-base origin/main HEAD)" - if ! git diff --diff-filter=ACM --quiet --exit-code "$MERGEBASE" -- '*.py' '*.pyi' &>/dev/null; then - git diff --name-only --diff-filter=ACM "$MERGEBASE" -- '*.py' '*.pyi' | xargs \ - codespell "${CODESPELL_EXCLUDES[@]}" - fi -} - -# Run Codespell -## This flag runs spell check of individual files. --files *must* be the first command line -## arg to use this option. -# if [[ "$1" == '--files' ]]; then -# spell_check "${@:2}" -# # If `--all` is passed, then any further arguments are ignored and the -# # entire python directory is linted. -# elif [[ "$1" == '--all' ]]; then -# spell_check_all -# else -# # Check spelling only of the files that changed in last commit. -# spell_check_changed -# fi -echo 'FastVideo codespell: Done' - - -# Lint specified files -lint() { - ruff check "$@" -} - -# Lint files that differ from main branch. Ignores dirs that are not slated -# for autolint yet. -lint_changed() { - # The `if` guard ensures that the list of filenames is not empty, which - # could cause ruff to receive 0 positional arguments, making it hang - # waiting for STDIN. - # - # `diff-filter=ACM` and $MERGEBASE is to ensure we only lint files that - # exist on both branches. - MERGEBASE="$(git merge-base origin/main HEAD)" - - if ! git diff --diff-filter=ACM --quiet --exit-code "$MERGEBASE" -- '*.py' '*.pyi' &>/dev/null; then - git diff --name-only --diff-filter=ACM "$MERGEBASE" -- '*.py' '*.pyi' | xargs \ - ruff check - fi - -} - -# Run Ruff -### This flag lints individual files. --files *must* be the first command line -### arg to use this option. -if [[ "$1" == '--files' ]]; then - lint "${@:2}" - # If `--all` is passed, then any further arguments are ignored and the - # entire python directory is linted. -elif [[ "$1" == '--all' ]]; then - lint fastvideo tests -else - # Format only the files that changed in last commit. - lint_changed -fi -echo 'FastVideo ruff: Done' - -# check spelling of specified files -isort_check() { - isort "$@" -} - -isort_check_all(){ - isort . -} - -# Spelling check of files that differ from main branch. -isort_check_changed() { - # The `if` guard ensures that the list of filenames is not empty, which - # could cause ruff to receive 0 positional arguments, making it hang - # waiting for STDIN. - # - # `diff-filter=ACM` and $MERGEBASE is to ensure we only lint files that - # exist on both branches. - MERGEBASE="$(git merge-base origin/main HEAD)" - - if ! git diff --diff-filter=ACM --quiet --exit-code "$MERGEBASE" -- '*.py' '*.pyi' &>/dev/null; then - git diff --name-only --diff-filter=ACM "$MERGEBASE" -- '*.py' '*.pyi' | xargs \ - isort - fi -} - -# Run Isort -# This flag runs spell check of individual files. --files *must* be the first command line -# arg to use this option. -if [[ "$1" == '--files' ]]; then - isort_check "${@:2}" - # If `--all` is passed, then any further arguments are ignored and the - # entire python directory is linted. -elif [[ "$1" == '--all' ]]; then - isort_check_all -else - # Check spelling only of the files that changed in last commit. - isort_check_changed -fi -echo 'FastVideo isort: Done' \ No newline at end of file diff --git a/pyproject.toml b/pyproject.toml index 6ff91550c..d5d714f3b 100644 --- a/pyproject.toml +++ b/pyproject.toml @@ -52,20 +52,7 @@ dependencies = [ # flash-attn: pip install flash-attn==2.7.0.post2 --no-build-isolation lint = [ - # formatting - "yapf==0.32.0", - "toml==0.10.2", - "tomli==2.0.2", - "ruff==0.6.5", - "codespell==2.3.0", - "isort==5.13.2", - "sphinx-lint==1.0.0", - - # type checking - "mypy==1.11.1", - "types-PyYAML", - "types-requests", - "types-setuptools", + "pre-commit==4.0.1", ] test = [